|
| 1 | +import abc |
1 | 2 | import traceback |
2 | 3 |
|
3 | 4 | import asyncio |
@@ -142,8 +143,21 @@ def get_device_ip(dev: Device) -> Optional[str]: |
142 | 143 | _logger.warning("get device ip error: %s", e) |
143 | 144 | return None |
144 | 145 |
|
| 146 | +class ApiMaker(metaclass=abc.ABCMeta): |
| 147 | + conf: AppSettings |
| 148 | + |
| 149 | + @asynccontextmanager |
| 150 | + async def make_api(self) -> AsyncIterator[Gnetcli]: |
| 151 | + async with GnetcliStarter(self.conf.server_path, self.conf.server_conf) as gnetcli_url: |
| 152 | + yield Gnetcli( |
| 153 | + server=gnetcli_url, |
| 154 | + auth_token=self.conf.make_server_credentials(), |
| 155 | + insecure_grpc=self.conf.insecure_grpc, |
| 156 | + user_agent="annet", |
| 157 | + ) |
| 158 | + |
145 | 159 |
|
146 | | -class GnetcliFetcher(Fetcher, AdapterWithConfig, AdapterWithName): |
| 160 | +class GnetcliFetcher(Fetcher, AdapterWithConfig, AdapterWithName, ApiMaker): |
147 | 161 | def __init__( |
148 | 162 | self, |
149 | 163 | url: Optional[str] = None, |
@@ -190,7 +204,7 @@ async def fetch( |
190 | 204 | ): |
191 | 205 | if not devices: |
192 | 206 | return {}, {} |
193 | | - async with make_api(self.conf) as api: |
| 207 | + async with self.make_api() as api: |
194 | 208 | return await self._fetch(api, devices, files_to_download, processes, max_slots) |
195 | 209 |
|
196 | 210 | async def _fetch( |
@@ -282,18 +296,7 @@ def parse_annet_qa(qa: list[annet.annlib.command.Question]) -> list[QA]: |
282 | 296 | return res |
283 | 297 |
|
284 | 298 |
|
285 | | -@asynccontextmanager |
286 | | -async def make_api(conf: AppSettings) -> AsyncIterator[Gnetcli]: |
287 | | - async with GnetcliStarter(conf.server_path, conf.server_conf) as gnetcli_url: |
288 | | - yield Gnetcli( |
289 | | - server=gnetcli_url, |
290 | | - auth_token= conf.make_server_credentials(), |
291 | | - insecure_grpc=conf.insecure_grpc, |
292 | | - user_agent="annet", |
293 | | - ) |
294 | | - |
295 | | - |
296 | | -class GnetcliDeployer(DeployDriver, AdapterWithConfig, AdapterWithName): |
| 299 | +class GnetcliDeployer(DeployDriver, AdapterWithConfig, AdapterWithName, ApiMaker): |
297 | 300 | def __init__( |
298 | 301 | self, |
299 | 302 | url: Optional[str] = None, |
@@ -335,7 +338,7 @@ async def bulk_deploy( |
335 | 338 | ) -> DeployResult: |
336 | 339 | if not deploy_cmds: |
337 | 340 | return DeployResult(hostnames=[], results={}, durations={}, original_states={}) |
338 | | - async with make_api(self.conf) as api: |
| 341 | + async with self.make_api() as api: |
339 | 342 | return await self._bulk_deploy( |
340 | 343 | api=api, |
341 | 344 | deploy_cmds=deploy_cmds, |
|
0 commit comments