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

33 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-05-06 10:21 +0200

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 

9from jumpstarter_protocol import jumpstarter_pb2_grpc 

10 

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 response = await jumpstarter_pb2_grpc.ExporterServiceStub(channel).GetReport(empty_pb2.Empty()) 

36 

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

38 topo[index] = [] 

39 

40 last_seen[report.uuid] = index 

41 

42 if report.parent_uuid != "": 

43 parent_index = last_seen[report.parent_uuid] 

44 topo[parent_index].append(index) 

45 

46 reports[index] = report 

47 

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

49 report = reports[index] 

50 

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

52 client = client_class( 

53 uuid=UUID(report.uuid), 

54 labels=report.labels, 

55 channel=channel, 

56 portal=portal, 

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

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

59 ) 

60 

61 clients[index] = client 

62 

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