Skip to content

Commit a9859d6

Browse files
jakob-kellerjakob-keller
committed
Generalize AioSession.warm_up_loader_caches() 
1 parent 78ef1ff commit a9859d6

2 files changed

Lines changed: 78 additions & 38 deletions

File tree

‎aiobotocore/session.py‎

Lines changed: 34 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import asyncio
2+
import contextlib
23
from typing import Optional
34

45
from botocore import UNSIGNED, translate
@@ -131,33 +132,46 @@ async def get_service_data(self, service_name, api_version=None):
131132

132133
def warm_up_loader_caches(
133134
self,
134-
service_name: str,
135+
service_name: str | None = None,
135136
api_version: Optional[str] = None,
136137
):
137138
loader = self.get_component('data_loader')
138139

139-
# from session.py
140+
# load generic data
141+
loader.load_data_with_path('_retry')
140142
loader.load_data_with_path('endpoints')
141-
loader.load_data('sdk-default-configuration')
142-
loader.load_service_model(service_name, 'waiters-2', api_version)
143-
loader.load_service_model(service_name, 'paginators-1', api_version)
144-
loader.load_service_model(
145-
service_name, type_name='service-2', api_version=api_version
146-
)
147-
loader.list_available_services(type_name='service-2')
143+
loader.load_data_with_path('partitions')
144+
loader.load_data_with_path('sdk-default-configuration')
145+
services = loader.list_available_services(type_name='service-2')
146+
147+
# load service-specific data
148+
for service_name in (service_name,) if service_name else services:
149+
# from session.py
150+
loader.load_service_model(
151+
service_name, type_name='service-2', api_version=api_version
152+
)
153+
with contextlib.suppress(UnknownServiceError):
154+
loader.load_service_model(
155+
service_name, 'paginators-1', api_version
156+
)
157+
with contextlib.suppress(UnknownServiceError):
158+
loader.load_service_model(
159+
service_name, 'waiters-2', api_version
160+
)
148161

149-
# from client.py
150-
loader.load_data('partitions')
151-
loader.load_service_model(
152-
service_name, 'service-2', api_version=api_version
153-
)
154-
loader.load_service_model(
155-
service_name, 'endpoint-rule-set-1', api_version=api_version
156-
)
157-
loader.load_data('_retry')
162+
# from client.py
163+
loader.load_service_model(
164+
service_name, 'service-2', api_version=api_version
165+
)
166+
loader.load_service_model(
167+
service_name, 'endpoint-rule-set-1', api_version=api_version
168+
)
158169

159-
# from docs/service.py
160-
loader.load_service_model(service_name, 'examples-1', api_version)
170+
# from docs/service.py
171+
with contextlib.suppress(UnknownServiceError):
172+
loader.load_service_model(
173+
service_name, 'examples-1', api_version
174+
)
161175

162176
def create_client(self, *args, **kwargs):
163177
return ClientCreatorContext(self._create_client(*args, **kwargs))

‎tests/test_session.py‎

Lines changed: 44 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import itertools
12
import logging
23

34
import pytest
@@ -59,6 +60,8 @@ async def test_set_user_agent_for_session(session: AioSession):
5960
@pytest.mark.parametrize(
6061
"service_name, api_version",
6162
[
63+
(None, None),
64+
("iot", None),
6265
("s3", None),
6366
("s3", "2006-03-01"),
6467
("ec2", "2016-11-16"),
@@ -67,35 +70,58 @@ async def test_set_user_agent_for_session(session: AioSession):
6770
def test_warm_up_loader_caches(
6871
session: AioSession, service_name, api_version, mocker
6972
):
73+
if service_name is None:
74+
services = [
75+
"ec2",
76+
"iot",
77+
"s3",
78+
]
79+
else:
80+
services = [service_name]
81+
7082
loader = mocker.Mock()
7183
get_component = mocker.patch.object(
7284
session, "get_component", return_value=loader
7385
)
86+
loader.list_available_services.return_value = services
7487

7588
session.warm_up_loader_caches(service_name, api_version)
7689

7790
get_component.assert_called_once_with("data_loader")
7891
assert loader.mock_calls == [
92+
# generic calls
93+
mocker.call.load_data_with_path("_retry"),
7994
mocker.call.load_data_with_path("endpoints"),
80-
mocker.call.load_data("sdk-default-configuration"),
81-
mocker.call.load_service_model(service_name, "waiters-2", api_version),
82-
mocker.call.load_service_model(
83-
service_name, "paginators-1", api_version
84-
),
85-
mocker.call.load_service_model(
86-
service_name, type_name="service-2", api_version=api_version
87-
),
95+
mocker.call.load_data_with_path("partitions"),
96+
mocker.call.load_data_with_path("sdk-default-configuration"),
8897
mocker.call.list_available_services(type_name="service-2"),
89-
mocker.call.load_data("partitions"),
90-
mocker.call.load_service_model(
91-
service_name, "service-2", api_version=api_version
92-
),
93-
mocker.call.load_service_model(
94-
service_name, "endpoint-rule-set-1", api_version=api_version
95-
),
96-
mocker.call.load_data("_retry"),
97-
mocker.call.load_service_model(
98-
service_name, "examples-1", api_version
98+
# service-specific calls
99+
*itertools.chain.from_iterable(
100+
(
101+
mocker.call.load_service_model(
102+
service_name,
103+
type_name="service-2",
104+
api_version=api_version,
105+
),
106+
mocker.call.load_service_model(
107+
service_name, "paginators-1", api_version
108+
),
109+
mocker.call.load_service_model(
110+
service_name, "waiters-2", api_version
111+
),
112+
mocker.call.load_service_model(
113+
service_name, "service-2", api_version=api_version
114+
),
115+
mocker.call.load_service_model(
116+
service_name,
117+
"endpoint-rule-set-1",
118+
api_version=api_version,
119+
),
120+
mocker.call.load_service_model(
121+
service_name, "examples-1", api_version
122+
),
123+
)
124+
for service_name in services
99125
),
100126
]
101127

0 commit comments

Comments
 (0)