def test_v2_client_v1_server(self): class Handler(BaseplateService.Iface): def is_healthy(self, context): return True handler = Handler() span_observer = mock.Mock(spec=SpanObserver) with serve_thrift(handler, BaseplateService) as server: with baseplate_thrift_client(server.endpoint, BaseplateServiceV2, span_observer) as context: healthy = context.example_service.is_healthy( request=IsHealthyRequest(probe=IsHealthyProbe.READINESS), ) self.assertTrue(healthy)
def test_v2_client_v2_server(self): class Handler(BaseplateServiceV2.Iface): def __init__(self): self.probe = None def is_healthy(self, context, req=None): self.probe = req.probe return True handler = Handler() span_observer = mock.Mock(spec=SpanObserver) with serve_thrift(handler, BaseplateServiceV2) as server: with baseplate_thrift_client(server.endpoint, BaseplateServiceV2, span_observer) as context: healthy = context.example_service.is_healthy( request=IsHealthyRequest(probe=IsHealthyProbe.LIVENESS), ) self.assertTrue(healthy) self.assertEqual(handler.probe, IsHealthyProbe.LIVENESS)
def check_thrift_service(endpoint: EndpointConfiguration, probe: int) -> None: pool = ThriftConnectionPool(endpoint, size=1, timeout=TIMEOUT) with pool.connection() as protocol: client = BaseplateService.Client(protocol) assert client.is_healthy(request=IsHealthyRequest( probe=probe), ), "service indicated unhealthiness in probe {probe}"