1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
|
from contextlib import AsyncExitStack, asynccontextmanager
from botocore.stub import Stubber
import aiobotocore.session
class StubbedSession(aiobotocore.session.AioSession):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._cached_clients = {}
self._client_stubs = {}
@asynccontextmanager
async def create_client(self, service_name, *args, **kwargs):
async with AsyncExitStack() as es:
es: AsyncExitStack
if service_name not in self._cached_clients:
client = await es.enter_async_context(
self._create_stubbed_client(service_name, *args, **kwargs)
)
self._cached_clients[service_name] = client
yield self._cached_clients[service_name]
@asynccontextmanager
async def _create_stubbed_client(self, service_name, *args, **kwargs):
async with AsyncExitStack() as es:
es: AsyncExitStack
client = await es.enter_async_context(
super().create_client(service_name, *args, **kwargs)
)
stubber = Stubber(client)
self._client_stubs[service_name] = stubber
yield client
@asynccontextmanager
async def stub(self, service_name, *args, **kwargs):
async with AsyncExitStack() as es:
es: AsyncExitStack
if service_name not in self._client_stubs:
await es.enter_async_context(
self.create_client(service_name, *args, **kwargs)
)
yield self._client_stubs[service_name]
def activate_stubs(self):
for stub in self._client_stubs.values():
stub.activate()
def verify_stubs(self):
for stub in self._client_stubs.values():
stub.assert_no_pending_responses()
|
