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