Coverage for /Users/ajo/work/jumpstarter/jumpstarter/packages/jumpstarter/jumpstarter/client/client.py: 100%

1from collections import OrderedDict, defaultdict 

2from contextlib import ExitStack, asynccontextmanager 

3from graphlib import TopologicalSorter 

4from uuid import UUID 

5 

6import grpc 

7from anyio.from_thread import BlockingPortal 

8from google.protobuf import empty_pb2 

9 

10from .grpc import MultipathExporterStub 

11from jumpstarter.client import DriverClient 

12from jumpstarter.common.importlib import import_class 

13 

14 

15@asynccontextmanager 

16async def client_from_path(path: str, portal: BlockingPortal, stack: ExitStack, allow: list[str], unsafe: bool): 

17 async with grpc.aio.secure_channel( 

18 f"unix://{path}", grpc.local_channel_credentials(grpc.LocalConnectionType.UDS) 

19 ) as channel: 

20 yield await client_from_channel(channel, portal, stack, allow, unsafe) 

21 

22 

23async def client_from_channel( 

24 channel: grpc.aio.Channel, 

25 portal: BlockingPortal, 

26 stack: ExitStack, 

27 allow: list[str], 

28 unsafe: bool, 

29) -> DriverClient: 

30 topo = defaultdict(list) 

31 last_seen = {} 

32 reports = {} 

33 clients = OrderedDict() 

34 

35 stub = MultipathExporterStub([channel]) 

36 

37 response = await stub.GetReport(empty_pb2.Empty()) 

38 

39 for index, report in enumerate(response.reports): 

40 topo[index] = [] 

41 

42 last_seen[report.uuid] = index 

43 

44 if report.parent_uuid != "": 

45 parent_index = last_seen[report.parent_uuid] 

46 topo[parent_index].append(index) 

47 

48 reports[index] = report 

49 

50 for index in TopologicalSorter(topo).static_order(): 

51 report = reports[index] 

52 

53 client_class = import_class(report.labels["jumpstarter.dev/client"], allow, unsafe) 

54 client = client_class( 

55 uuid=UUID(report.uuid), 

56 labels=report.labels, 

57 stub=stub, 

58 portal=portal, 

59 stack=stack.enter_context(ExitStack()), 

60 children={reports[k].labels["jumpstarter.dev/name"]: clients[k] for k in topo[index]}, 

61 ) 

62 

63 clients[index] = client 

64 

65 return clients.popitem(last=True)[1]