1# ------------------------------------ 2# Copyright (c) Microsoft Corporation. 3# Licensed under the MIT License. 4# ------------------------------------ 5import asyncio 6 7from azure_devtools.scenario_tests.patches import mock_in_unit_test 8from devtools_testutils import AzureTestCase 9 10 11def skip_sleep(unit_test): 12 async def immediate_return(_): 13 return 14 15 return mock_in_unit_test(unit_test, "asyncio.sleep", immediate_return) 16 17 18class KeyVaultTestCase(AzureTestCase): 19 def __init__(self, *args, match_body=True, **kwargs): 20 super().__init__(*args, match_body=match_body, **kwargs) 21 self.replay_patches.append(skip_sleep) 22 23 def setUp(self): 24 self.list_test_size = 7 25 super(KeyVaultTestCase, self).setUp() 26 27 def get_resource_name(self, name): 28 """helper to create resources with a consistent, test-indicative prefix""" 29 return super(KeyVaultTestCase, self).get_resource_name("livekvtest{}".format(name)) 30 31 async def _poll_until_no_exception(self, fn, expected_exception, max_retries=20, retry_delay=3): 32 """polling helper for live tests because some operations take an unpredictable amount of time to complete""" 33 34 for i in range(max_retries): 35 try: 36 return await fn() 37 except expected_exception: 38 if i == max_retries - 1: 39 raise 40 if self.is_live: 41 await asyncio.sleep(retry_delay) 42 43 async def _poll_until_exception(self, fn, expected_exception, max_retries=20, retry_delay=3): 44 """polling helper for live tests because some operations take an unpredictable amount of time to complete""" 45 46 for _ in range(max_retries): 47 try: 48 await fn() 49 if self.is_live: 50 await asyncio.sleep(retry_delay) 51 except expected_exception: 52 return 53 self.fail("expected exception {expected_exception} was not raised") 54