def baseplate_thrift_client(endpoint, client_span_observer=None): pool = ThriftConnectionPool(endpoint) context = RequestContext() server_span = ServerSpan( trace_id=1234, parent_id=2345, span_id=3456, flags=4567, sampled=1, name="example_service.example", context=context, ) if client_span_observer: class TestServerSpanObserver(ServerSpanObserver): def on_child_span_created(self, span): span.register(client_span_observer) server_span.register(TestServerSpanObserver()) edge_context_factory = make_edge_context_factory() edge_context = edge_context_factory.from_upstream( SERIALIZED_EDGECONTEXT_WITH_VALID_AUTH) edge_context.attach_context(context) context_factory = ThriftContextFactory(pool, TestService.Client) client = context_factory.make_object_for_context("example_service", server_span) setattr(context, "example_service", client) yield context
def test_client_proxy_flow(self): client_memory_trans = TMemoryBuffer() client_prot = THeaderProtocol(client_memory_trans) class Pool(object): @contextlib.contextmanager def connection(self): yield client_prot client_factory = ThriftContextFactory(Pool(), TestService.Client) span = mock.MagicMock() child_span = span.make_child().__enter__() child_span.trace_id = 1 child_span.parent_id = 1 child_span.id = 1 child_span.sampled = True child_span.flags = None # We decode the token to unicode to make sure that it is converted to # bytes correctly by the AuthenticationContext. We do this because a # unicode token in Python 2 ends up causing a UnicodeDecodeError when # Thrift tries to write the header. unicode_token = self.VALID_TOKEN.decode() auth_context = AuthenticationContext( token=unicode_token, secrets=self.secrets, ) edge_context = EdgeRequestContext( authentication_context=auth_context, header=self.SERIALIZED_REQUEST_HEADER, ) edge_context.attach_context(child_span.context) client = client_factory.make_object_for_context("test", span) try: client.example_simple() except TTransportException: 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) context, _ = self.observer.on_server_span_created.call_args[0] try: self.assertEqual(context.request_context.user.id, "test_user_id") self.assertEqual(context.request_context.user.roles, set()) self.assertEqual(context.request_context.user.is_logged_in, True) self.assertEqual(context.request_context.user.loid, "t2_deadbeef") self.assertEqual(context.request_context.user.cookie_created_ms, 100000) self.assertEqual(context.request_context.oauth_client.id, None) self.assertFalse( context.request_context.oauth_client.is_type("third_party")) self.assertEqual(context.request_context.session.id, "beefdead") except jwt.exceptions.InvalidAlgorithmError: raise unittest.SkipTest("cryptography is not installed")
def test_client_proxy_flow(self): client_memory_trans = TMemoryBuffer() client_prot = THeaderProtocolFactory().getProtocol(client_memory_trans) class Pool(object): @contextlib.contextmanager def connection(self): yield client_prot client_factory = ThriftContextFactory(Pool(), TestService.Client) span = mock.MagicMock() child_span = span.make_child().__enter__() child_span.trace_id = 1 child_span.parent_id = 1 child_span.id = 1 child_span.sampled = True child_span.flags = None edge_context = self.edge_context_factory.from_upstream( SERIALIZED_EDGECONTEXT_WITH_VALID_AUTH) edge_context.attach_context(child_span.context) client = client_factory.make_object_for_context("test", span) try: client.example_simple() except (TTransportException, EOFError): pass # we don't have a test response for the client itrans = TMemoryBuffer(client_memory_trans.getvalue()) iprot = THeaderProtocolFactory().getProtocol(itrans) self.processor.process(iprot, self.oprot) context, _ = self.observer.on_server_span_created.call_args[0] try: self.assertEqual(context.request_context.user.id, "t2_example") self.assertEqual(context.request_context.user.roles, set()) self.assertEqual(context.request_context.user.is_logged_in, True) self.assertEqual(context.request_context.user.loid, "t2_deadbeef") self.assertEqual(context.request_context.user.cookie_created_ms, 100000) self.assertEqual(context.request_context.oauth_client.id, None) self.assertFalse( context.request_context.oauth_client.is_type("third_party")) self.assertEqual(context.request_context.session.id, "beefdead") except jwt.exceptions.InvalidAlgorithmError: raise unittest.SkipTest("cryptography is not installed")