def test_tween_is_wrapped_by_zipkin_logging_context( binann_mock, endp_mock, attrs_mock, context): instance = context.return_value attrs_mock.return_value = 'zipkin_attrs' binann_mock.return_value = {'k': 'v'} endp_mock.return_value = 'endpoint_attrs' tween = zipkin.zipkin_tween(mock.Mock(), 'registry') tween(mock.Mock()) instance.__enter__.assert_called_once_with() instance.__exit__.assert_called_once_with(None, None, None)
def _test_zipkin_context( binann_mock, endp_mock, attrs_mock, context, is_sampled): instance = context.return_value attrs_mock.return_value = ZipkinAttrs( 'trace_id', 'span_id', 'parent_span_id', 'flags', is_sampled) binann_mock.return_value = {'k': 'v'} endp_mock.return_value = 'thrift_endpoint' tween = zipkin.zipkin_tween(mock.Mock(), 'registry') tween(mock.Mock()) return instance
def test_tween_not_sampled_sets_zipkin_trace_id(is_tracing_mock): # Tests that even in the unsampled case, the zipkin tween sets # the zipkin_trace_id attr on the request. is_tracing_mock.return_value = False assert get_zipkin_attrs() is None request = Request.blank('/', headers={'X-B3-TraceId': 'deadbeefdeadbeef'}) tween = zipkin.zipkin_tween(mock.Mock(), 'registry') tween(request) assert request.zipkin_trace_id == 'deadbeefdeadbeef' # Make sure the tween doesn't leave zipkin attrs on threadlocal storage assert get_zipkin_attrs() is None