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')