Coverage for /Users/ajo/work/jumpstarter/jumpstarter/packages/jumpstarter/jumpstarter/common/streams.py: 75%

28 statements  

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

1from contextlib import asynccontextmanager 

2from typing import Annotated, Literal, Union 

3from uuid import UUID 

4 

5import grpc 

6from jumpstarter_protocol.jumpstarter.v1 import router_pb2_grpc 

7from pydantic import BaseModel, Field, Json 

8 

9from jumpstarter.common.grpc import aio_secure_channel, ssl_channel_credentials 

10from jumpstarter.streams.common import forward_stream 

11from jumpstarter.streams.router import RouterStream 

12 

13 

14class ResourceStreamRequest(BaseModel): 

15 kind: Literal["resource"] = "resource" 

16 uuid: UUID 

17 

18 

19class DriverStreamRequest(BaseModel): 

20 kind: Literal["driver"] = "driver" 

21 uuid: UUID 

22 method: str 

23 

24 

25StreamRequest = Annotated[ 

26 Union[ResourceStreamRequest, DriverStreamRequest], 

27 Field(discriminator="kind"), 

28] 

29 

30 

31class StreamRequestMetadata(BaseModel): 

32 request: Json[StreamRequest] 

33 

34 

35@asynccontextmanager 

36async def connect_router_stream(endpoint, token, stream, tls_config, grpc_options): 

37 credentials = grpc.composite_channel_credentials( 

38 await ssl_channel_credentials(endpoint, tls_config), 

39 grpc.access_token_call_credentials(token), 

40 ) 

41 

42 async with aio_secure_channel(endpoint, credentials, grpc_options) as channel: 

43 router = router_pb2_grpc.RouterServiceStub(channel) 

44 context = router.Stream(metadata=()) 

45 async with RouterStream(context=context) as s: 

46 async with forward_stream(s, stream): 

47 yield