Skip to content

Commit 02280fd

Browse files
authored
feat: allow to override the api creation function in gnetcli fetcher and deployer (#76)
1 parent cc879e7 commit 02280fd

1 file changed

Lines changed: 18 additions & 15 deletions

File tree

src/gnetcli_adapter/gnetcli_adapter.py

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import abc
12
import traceback
23

34
import asyncio
@@ -142,8 +143,21 @@ def get_device_ip(dev: Device) -> Optional[str]:
142143
_logger.warning("get device ip error: %s", e)
143144
return None
144145

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+
145159

146-
class GnetcliFetcher(Fetcher, AdapterWithConfig, AdapterWithName):
160+
class GnetcliFetcher(Fetcher, AdapterWithConfig, AdapterWithName, ApiMaker):
147161
def __init__(
148162
self,
149163
url: Optional[str] = None,
@@ -190,7 +204,7 @@ async def fetch(
190204
):
191205
if not devices:
192206
return {}, {}
193-
async with make_api(self.conf) as api:
207+
async with self.make_api() as api:
194208
return await self._fetch(api, devices, files_to_download, processes, max_slots)
195209

196210
async def _fetch(
@@ -282,18 +296,7 @@ def parse_annet_qa(qa: list[annet.annlib.command.Question]) -> list[QA]:
282296
return res
283297

284298

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):
297300
def __init__(
298301
self,
299302
url: Optional[str] = None,
@@ -335,7 +338,7 @@ async def bulk_deploy(
335338
) -> DeployResult:
336339
if not deploy_cmds:
337340
return DeployResult(hostnames=[], results={}, durations={}, original_states={})
338-
async with make_api(self.conf) as api:
341+
async with self.make_api() as api:
339342
return await self._bulk_deploy(
340343
api=api,
341344
deploy_cmds=deploy_cmds,

0 commit comments

Comments
 (0)