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
