def test_with_headers(self): client_memory_trans = TMemoryBuffer() client_prot = THeaderProtocol(client_memory_trans) client_header_trans = client_prot.trans client_header_trans.set_header("Trace", "1234") client_header_trans.set_header("Parent", "2345") client_header_trans.set_header("Span", "3456") client_header_trans.set_header("Sampled", "1") client_header_trans.set_header("Flags", "1") client = BaseplateService.Client(client_prot) try: client.is_healthy() except: pass # we don't have a test response for the client self.itrans._readBuffer = StringIO(client_memory_trans.getvalue()) self.processor.process(self.iprot, self.oprot, self.server_context) self.assertEqual(self.observer.on_server_span_created.call_count, 1) context, server_span = self.observer.on_server_span_created.call_args[ 0] self.assertEqual(server_span.trace_id, 1234) self.assertEqual(server_span.parent_id, 2345) self.assertEqual(server_span.id, 3456) self.assertTrue(server_span.sampled) self.assertEqual(server_span.flags, 1) self.assertTrue(self.server_observer.on_start.called) self.assertTrue(self.server_observer.on_finish.called)
def setUp(self): self.itrans = TMemoryBuffer() self.iprot = THeaderProtocol(self.itrans) self.otrans = TMemoryBuffer() self.oprot = THeaderProtocol(self.otrans) self.observer = mock.Mock(spec=BaseplateObserver) self.server_observer = mock.Mock(spec=ServerSpanObserver) def _register_mock(context, server_span): server_span.register(self.server_observer) self.observer.on_server_span_created.side_effect = _register_mock self.logger = mock.Mock(spec=logging.Logger) self.server_context = TRpcConnectionContext(self.itrans, self.iprot, self.oprot) baseplate = Baseplate() baseplate.register(self.observer) event_handler = BaseplateProcessorEventHandler(self.logger, baseplate) handler = TestHandler() self.processor = BaseplateService.ContextProcessor(handler) self.processor.setEventHandler(event_handler)
def test_no_headers(self): client_memory_trans = TMemoryBuffer() client_prot = THeaderProtocol(client_memory_trans) client = BaseplateService.Client(client_prot) try: client.is_healthy() except: pass # we don't have a test response for the client self.itrans._readBuffer = StringIO(client_memory_trans.getvalue()) self.processor.process(self.iprot, self.oprot, self.server_context) self.assertEqual(self.observer.on_root_span_created.call_count, 1) context, root_span = self.observer.on_root_span_created.call_args[0] self.assertEqual(root_span.trace_id, "no-trace") self.assertEqual(root_span.parent_id, "no-parent") self.assertEqual(root_span.id, "no-span") self.assertTrue(self.root_observer.on_start.called) self.assertTrue(self.root_observer.on_stop.called)
def test_no_trace_headers(self, getrandbits): getrandbits.return_value = 1234 client_memory_trans = TMemoryBuffer() client_prot = THeaderProtocol(client_memory_trans) client = BaseplateService.Client(client_prot) try: client.is_healthy() except: pass # we don't have a test response for the client self.itrans._readBuffer = StringIO(client_memory_trans.getvalue()) self.processor.process(self.iprot, self.oprot, self.server_context) self.assertEqual(self.observer.on_server_span_created.call_count, 1) context, server_span = self.observer.on_server_span_created.call_args[0] self.assertEqual(server_span.trace_id, 1234) self.assertEqual(server_span.parent_id, None) self.assertEqual(server_span.id, 1234) self.assertTrue(self.server_observer.on_start.called) self.assertTrue(self.server_observer.on_finish.called)
def check_thrift_service(endpoint: EndpointConfiguration) -> None: pool = ThriftConnectionPool(endpoint, size=1, timeout=TIMEOUT) with pool.connection() as protocol: client = BaseplateService.Client(protocol) assert client.is_healthy(), "service indicated unhealthiness"
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}"
def check_thrift_service(endpoint): protocol = _make_protocol(endpoint) protocol.trans.getTransport().setTimeout(TIMEOUT * 1000.) protocol.trans.open() client = BaseplateService.Client(protocol) assert client.is_healthy(), "service indicated unhealthiness"