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"
