def test_proxy_methods_attached(self, mock_enumerate): mock_enumerate.return_value = ["one", "two"] proxy = thrift.PooledClientProxy(self.mock_client_cls, self.mock_pool, self.mock_server_span, "namespace") self.assertTrue(callable(proxy.one)) self.assertTrue(callable(proxy.two))
def test_call_proxy_method(self, mock_enumerate): mock_enumerate.return_value = ["one", "two"] proxy = thrift.PooledClientProxy( self.mock_client_cls, self.mock_pool, self.mock_server_span, "namespace") result = proxy.one(mock.sentinel.first, mock.sentinel.second) self.assertEqual(self.mock_client.one.call_count, 1) self.assertEqual(result, self.mock_client.one.return_value) self.assertEqual(self.mock_server_span.make_child.call_args, mock.call("namespace.one"))
def test_null_edge_request_headers_not_set(self, mock_enumerate): mock_enumerate.return_value = ["one", "two"] child_span = self.mock_server_span.make_child() child_span.context.raw_request_context = None proxy = thrift.PooledClientProxy(self.mock_client_cls, self.mock_pool, self.mock_server_span, "namespace") proxy.one(mock.sentinel.first, mock.sentinel.second) self.assertNotIn(b'Edge-Request', self.outbound_headers)
def test_edge_request_headers(self, mock_enumerate): mock_enumerate.return_value = ["one", "two"] child_span = self.mock_server_span.make_child().__enter__() child_span.context.raw_request_context = "edge_request_context" proxy = thrift.PooledClientProxy( self.mock_client_cls, self.mock_pool, self.mock_server_span, "namespace") proxy.one(mock.sentinel.first, mock.sentinel.second) self.assertEqual(self.outbound_headers.get(b'Edge-Request'), 'edge_request_context')
def test_null_edge_request_headers_not_set(self, mock_enumerate): mock_enumerate.return_value = ["one", "two"] headers = {} def mock_set_header(header, val): headers[header] = val self.initialize_edge_request_headers(mock_set_header, None, None) proxy = thrift.PooledClientProxy(self.mock_client_cls, self.mock_pool, self.mock_server_span, "namespace") result = proxy.one(mock.sentinel.first, mock.sentinel.second) self.assertNotIn('Authentication', headers) self.assertNotIn('Edge-Request', headers)
def test_edge_request_headers(self, mock_enumerate): mock_enumerate.return_value = ["one", "two"] headers = {} def mock_set_header(header, val): headers[header] = val self.initialize_edge_request_headers(mock_set_header, 'authn_token', 'edge_request_context') proxy = thrift.PooledClientProxy(self.mock_client_cls, self.mock_pool, self.mock_server_span, "namespace") result = proxy.one(mock.sentinel.first, mock.sentinel.second) self.assertEqual(headers['Authentication'], 'authn_token') self.assertEqual(headers['Edge-Request'], 'edge_request_context')