def test_init(self, mock_generate_kind):
# Test that all arguments are correctly saved
zipkin_attrs = ZipkinAttrs(None, None, None, None, None)
transport = MockTransportHandler()
firehose = MockTransportHandler()
stack = Stack([])
span_storage = SpanStorage()
tracer = MockTracer()
context = tracer.zipkin_span(
service_name="test_service",
span_name="test_span",
zipkin_attrs=zipkin_attrs,
transport_handler=transport,
max_span_batch_size=10,
annotations={"test_annotation": 1},
binary_annotations={"status": "200"},
port=80,
sample_rate=100.0,
include=("cs", "cr"),
add_logging_annotation=True,
report_root_timestamp=True,
use_128bit_trace_id=True,
host="127.0.0.255",
context_stack=stack,
span_storage=span_storage,
firehose_handler=firehose,
kind=Kind.CLIENT,
timestamp=1234,
duration=10,
encoding=Encoding.V2_JSON,
)
assert context.service_name == "test_service"
assert context.span_name == "test_span"
assert context.zipkin_attrs_override == zipkin_attrs
assert context.transport_handler == transport
assert context.max_span_batch_size == 10
assert context.annotations == {"test_annotation": 1}
assert context.binary_annotations == {"status": "200"}
assert context.port == 80
assert context.sample_rate == 100.0
assert context.add_logging_annotation is True
assert context.report_root_timestamp_override is True
assert context.use_128bit_trace_id is True
assert context.host == "127.0.0.255"
assert context._context_stack == stack
assert context._span_storage == span_storage
assert context.firehose_handler == firehose
assert mock_generate_kind.call_count == 1
assert mock_generate_kind.call_args == mock.call(
context, Kind.CLIENT, ("cs", "cr"),
)
assert context.timestamp == 1234
assert context.duration == 10
assert context.encoding == Encoding.V2_JSON
assert context._tracer == tracer
# Check for backward compatibility
assert tracer.get_spans() == span_storage
assert tracer.get_context() == stack
def test_stop_with_bad_error_that_cannot_be_stringified(self):
class BadExceptionThatCannotBeStringified(Exception):
def __str__(self):
# Returning a non-str will cause a TypeError.
return 42
__unicode__ = __str__
# Transport is not setup, exit immediately
transport = MockTransportHandler()
span_storage = SpanStorage()
context = zipkin.zipkin_span(
service_name="test_service",
span_name="test_span",
transport_handler=transport,
sample_rate=100.0,
span_storage=span_storage,
)
with mock.patch.object(context,
"update_binary_annotations") as mock_upd:
context.start()
context.stop(
BadExceptionThatCannotBeStringified,
BadExceptionThatCannotBeStringified(),
)
assert mock_upd.call_args == mock.call({
zipkin.ERROR_KEY:
"BadExceptionThatCannotBeStringified: BadExceptionThatCannotBeStringified()" # noqa
})
assert len(span_storage) == 0
def test_zipkin_span_trace_with_no_sampling(
logging_context_cls_mock,
create_endpoint_mock,
create_attrs_for_span_mock,
mock_context_stack,
):
zipkin_attrs = ZipkinAttrs(
trace_id='0',
span_id='1',
parent_span_id=None,
flags='0',
is_sampled=False,
)
with zipkin.zipkin_span(
service_name='my_service',
span_name='span_name',
zipkin_attrs=zipkin_attrs,
transport_handler=MockTransportHandler(),
port=5,
context_stack=mock_context_stack,
span_storage=SpanStorage(),
):
pass
assert create_attrs_for_span_mock.call_count == 0
mock_context_stack.push.assert_called_once_with(zipkin_attrs, )
assert create_endpoint_mock.call_count == 0
assert logging_context_cls_mock.call_count == 0
mock_context_stack.pop.assert_called_once_with()
def test_batch_sender_add_span_not_sampled_with_firehose(
add_span_mock, flush_mock, time_mock):
attr = ZipkinAttrs(
trace_id='0000000000000001',
span_id='0000000000000002',
parent_span_id=None,
flags=None,
is_sampled=False,
)
span_storage = SpanStorage()
transport_handler = mock.Mock()
firehose_handler = mock.Mock()
context = logging_helper.ZipkinLoggingContext(
zipkin_attrs=attr,
endpoint=create_endpoint(80, 'test_server', '127.0.0.1'),
span_name='span_name',
transport_handler=transport_handler,
report_root_timestamp=False,
span_storage=span_storage,
firehose_handler=firehose_handler,
service_name='test_server',
encoding=Encoding.V1_JSON,
)
context.start_timestamp = 24
context.response_status_code = 200
context.binary_annotations_dict = {'k': 'v'}
time_mock.return_value = 42
context.emit_spans()
assert add_span_mock.call_count == 1
assert flush_mock.call_count == 1
def test_batch_sender_add_span_not_called_if_not_sampled(
add_span_mock, flush_mock):
attr = ZipkinAttrs(
trace_id='0000000000000001',
span_id='0000000000000002',
parent_span_id=None,
flags=None,
is_sampled=False,
)
span_storage = SpanStorage()
transport_handler = mock.Mock()
context = logging_helper.ZipkinLoggingContext(
zipkin_attrs=attr,
endpoint=create_endpoint(80, 'test_server', '127.0.0.1'),
span_name='span_name',
transport_handler=transport_handler,
report_root_timestamp=False,
span_storage=span_storage,
service_name='test_server',
encoding=Encoding.V1_JSON,
)
context.emit_spans()
assert add_span_mock.call_count == 0
assert flush_mock.call_count == 0
def test_zipkin_logging_client_context_emit_spans(add_span_mock, flush_mock,
time_mock, fake_endpoint):
# This lengthy function tests that the logging context properly
# logs root client span
trace_id = '000000000000000f'
client_span_id = '0000000000000003'
attr = ZipkinAttrs(
trace_id=trace_id,
span_id=client_span_id,
parent_span_id=None,
flags=None,
is_sampled=True,
)
span_storage = SpanStorage()
transport_handler = mock.Mock()
context = logging_helper.ZipkinLoggingContext(
zipkin_attrs=attr,
endpoint=fake_endpoint,
span_name='GET /foo',
transport_handler=transport_handler,
report_root_timestamp=True,
span_storage=span_storage,
client_context=True,
service_name='test_server',
encoding=Encoding.V1_JSON,
)
context.start_timestamp = 24
context.response_status_code = 200
context.binary_annotations_dict = {'k': 'v'}
time_mock.return_value = 42
context.emit_spans()
log_call = add_span_mock.call_args_list[0]
assert log_call[0][1].build_v1_span() == SpanBuilder(
trace_id=trace_id,
name='GET /foo',
parent_id=None,
span_id=client_span_id,
timestamp=24.0,
duration=18.0,
annotations={
'cs': 24,
'cr': 42
},
tags={
'k': 'v'
},
kind=Kind.CLIENT,
service_name='test_server',
local_endpoint=fake_endpoint,
).build_v1_span()
assert flush_mock.call_count == 1
def test_init_defaults(self, mock_generate_kind, mock_storage):
# Test that special arguments are properly defaulted
mock_storage.return_value = SpanStorage()
context = zipkin.zipkin_span(
service_name="test_service", span_name="test_span",
)
assert context.service_name == "test_service"
assert context.span_name == "test_span"
assert context.annotations == {}
assert context.binary_annotations == {}
assert mock_generate_kind.call_args == mock.call(context, None, None)
def context():
attr = ZipkinAttrs(None, None, None, None, False)
return logging_helper.ZipkinLoggingContext(
zipkin_attrs=attr,
endpoint=create_endpoint(80, 'test_server', '127.0.0.1'),
span_name='span_name',
transport_handler=MockTransportHandler(),
report_root_timestamp=False,
span_storage=SpanStorage(),
service_name='test_server',
encoding=Encoding.V1_JSON,
)
def test_zipkin_span_passed_sampled_attrs(
logging_context_cls_mock,
create_endpoint_mock,
create_attrs_for_span_mock,
mock_context_stack,
):
# Make sure that if zipkin_span is passed *sampled* ZipkinAttrs, but is
# also configured to do sampling itself, the passed ZipkinAttrs are used.
transport_handler = MockTransportHandler()
zipkin_attrs = ZipkinAttrs(
trace_id='0',
span_id='1',
parent_span_id=None,
flags='0',
is_sampled=True,
)
span_storage = SpanStorage()
with zipkin.zipkin_span(
service_name='some_service_name',
span_name='span_name',
transport_handler=transport_handler,
port=5,
sample_rate=100.0,
zipkin_attrs=zipkin_attrs,
context_stack=mock_context_stack,
span_storage=span_storage,
) as zipkin_context:
assert zipkin_context.port == 5
assert not create_attrs_for_span_mock.called
mock_context_stack.push.assert_called_once_with(zipkin_attrs)
create_endpoint_mock.assert_called_once_with(5, 'some_service_name', None)
# Logging context should not report timestamp/duration for the server span,
# since it's assumed that the client part of this span will do that.
logging_context_cls_mock.assert_called_once_with(
zipkin_attrs,
create_endpoint_mock.return_value,
'span_name',
transport_handler,
False,
span_storage,
'some_service_name',
binary_annotations={},
add_logging_annotation=False,
client_context=False,
max_span_batch_size=None,
firehose_handler=None,
encoding=Encoding.V1_THRIFT,
)
mock_context_stack.pop.assert_called_once_with()
def test_zipkin_client_span_decorator(
logging_context_cls_mock,
create_endpoint_mock,
create_attrs_for_span_mock,
mock_context_stack,
):
transport_handler = MockTransportHandler()
span_storage = SpanStorage()
@zipkin.zipkin_span(
service_name='some_service_name',
span_name='span_name',
transport_handler=transport_handler,
port=5,
sample_rate=100.0,
include=('client', ),
host='1.5.1.2',
context_stack=mock_context_stack,
span_storage=span_storage,
)
def test_func(a, b):
return a + b
assert test_func(1, 2) == 3
create_attrs_for_span_mock.assert_called_once_with(
sample_rate=100.0,
use_128bit_trace_id=False,
)
mock_context_stack.push.assert_called_once_with(
create_attrs_for_span_mock.return_value, )
create_endpoint_mock.assert_called_once_with(5, 'some_service_name',
'1.5.1.2')
# The decorator was passed a sample rate and no Zipkin attrs, so it's
# assumed to be the root of a trace and it should report timestamp/duration
assert logging_context_cls_mock.call_args == mock.call(
create_attrs_for_span_mock.return_value,
create_endpoint_mock.return_value,
'span_name',
transport_handler,
True,
span_storage,
'some_service_name',
binary_annotations={},
add_logging_annotation=False,
client_context=True,
max_span_batch_size=None,
firehose_handler=None,
encoding=Encoding.V1_THRIFT,
)
mock_context_stack.pop.assert_called_once_with()
def test_stop_no_transport(self):
# Transport is not setup, exit immediately
span_storage = SpanStorage()
context = zipkin.zipkin_span(
service_name='test_service',
span_name='test_span',
span_storage=span_storage,
)
with mock.patch.object(context, 'logging_context') as mock_log_ctx:
context.start()
context.stop()
assert mock_log_ctx.stop.call_count == 0
assert len(span_storage) == 0
def get_thread_local_span_storage():
"""A wrapper to return _thread_local.span_storage
Returns a SpanStorage object used to temporarily store all spans created in
the current process. The transport handlers will pull from this storage when
they emit the spans.
:returns: SpanStore object containing all non-root spans.
:rtype: py_zipkin.storage.SpanStore
"""
if not hasattr(_thread_local, 'span_storage'):
from py_zipkin.storage import SpanStorage
_thread_local.span_storage = SpanStorage()
return _thread_local.span_storage
def test_zipkin_trace_context_attrs_is_always_popped(
logging_context_cls_mock,
create_endpoint_mock,
create_attrs_for_span_mock,
mock_context_stack,
):
with pytest.raises(Exception):
with zipkin.zipkin_span(service_name='my_service',
span_name='my_span_name',
transport_handler=MockTransportHandler(),
port=22,
sample_rate=100.0,
context_stack=mock_context_stack,
span_storage=SpanStorage()):
raise Exception
mock_context_stack.pop.assert_called_once_with()
def test_zipkin_span_for_new_trace(
logging_context_cls_mock,
create_endpoint_mock,
create_attrs_for_span_mock,
mock_context_stack,
):
transport_handler = MockTransportHandler()
firehose_handler = mock.Mock()
span_storage = SpanStorage()
with zipkin.zipkin_span(
service_name='some_service_name',
span_name='span_name',
transport_handler=transport_handler,
port=5,
sample_rate=100.0,
context_stack=mock_context_stack,
span_storage=span_storage,
firehose_handler=firehose_handler,
) as zipkin_context:
assert zipkin_context.port == 5
pass
create_attrs_for_span_mock.assert_called_once_with(
sample_rate=100.0,
use_128bit_trace_id=False,
)
mock_context_stack.push.assert_called_once_with(
create_attrs_for_span_mock.return_value, )
create_endpoint_mock.assert_called_once_with(5, 'some_service_name', None)
assert logging_context_cls_mock.call_args == mock.call(
create_attrs_for_span_mock.return_value,
create_endpoint_mock.return_value,
'span_name',
transport_handler,
True,
span_storage,
'some_service_name',
binary_annotations={},
add_logging_annotation=False,
client_context=False,
max_span_batch_size=None,
firehose_handler=firehose_handler,
encoding=Encoding.V1_THRIFT,
)
mock_context_stack.pop.assert_called_once_with()
def test_stop_with_error(self):
# Transport is not setup, exit immediately
transport = MockTransportHandler()
span_storage = SpanStorage()
context = zipkin.zipkin_span(
service_name='test_service',
span_name='test_span',
transport_handler=transport,
sample_rate=100.0,
span_storage=span_storage,
)
with mock.patch.object(context, 'update_binary_annotations') as mock_upd:
context.start()
context.stop(ValueError, 'bad error')
assert mock_upd.call_args == mock.call({
zipkin.ERROR_KEY: 'ValueError: bad error'
})
assert len(span_storage) == 0
def test_zipkin_span_trace_with_0_sample_rate(
logging_context_cls_mock,
create_endpoint_mock,
create_attrs_for_span_mock,
mock_context_stack,
firehose_enabled,
):
create_attrs_for_span_mock.return_value = ZipkinAttrs(
trace_id=generate_random_64bit_string(),
span_id=generate_random_64bit_string(),
parent_span_id=None,
flags='0',
is_sampled=False,
)
transport_handler = MockTransportHandler()
span_storage = SpanStorage()
with zipkin.zipkin_span(service_name='some_service_name',
span_name='span_name',
transport_handler=transport_handler,
sample_rate=0.0,
context_stack=mock_context_stack,
span_storage=span_storage,
firehose_handler=mock.Mock()
if firehose_enabled else None) as zipkin_context:
assert zipkin_context.port == 0
create_attrs_for_span_mock.assert_called_once_with(
sample_rate=0.0,
use_128bit_trace_id=False,
)
mock_context_stack.push.assert_called_once_with(
create_attrs_for_span_mock.return_value)
# When firehose mode is on, we log regardless of sample rate
assert create_endpoint_mock.call_count == (1 if firehose_enabled else 0)
assert logging_context_cls_mock.call_count == (1
if firehose_enabled else 0)
mock_context_stack.pop.assert_called_once_with()
def test_zipkin_logging_server_context_emit_spans(add_span_mock, flush_mock,
time_mock, fake_endpoint):
# This lengthy function tests that the logging context properly
# logs both client and server spans.
trace_id = '000000000000000f'
parent_span_id = '0000000000000001'
server_span_id = '0000000000000002'
client_span_id = '0000000000000003'
client_span_name = 'breadcrumbs'
client_svc_name = 'svc'
attr = ZipkinAttrs(
trace_id=trace_id,
span_id=server_span_id,
parent_span_id=parent_span_id,
flags=None,
is_sampled=True,
)
span_storage = SpanStorage()
client_span = SpanBuilder(
trace_id=trace_id,
name=client_span_name,
parent_id=server_span_id,
span_id=client_span_id,
timestamp=26.0,
duration=4.0,
annotations={
'ann2': 2,
'cs': 26,
'cr': 30
},
tags={'bann2': 'yiss'},
kind=Kind.CLIENT,
service_name=client_svc_name,
)
span_storage.append(client_span)
transport_handler = mock.Mock()
context = logging_helper.ZipkinLoggingContext(
zipkin_attrs=attr,
endpoint=fake_endpoint,
span_name='GET /foo',
transport_handler=transport_handler,
report_root_timestamp=True,
span_storage=span_storage,
service_name='test_server',
encoding=Encoding.V1_JSON,
)
context.start_timestamp = 24
context.response_status_code = 200
context.binary_annotations_dict = {'k': 'v'}
time_mock.return_value = 42
context.emit_spans()
client_log_call, server_log_call = add_span_mock.call_args_list
assert server_log_call[0][1].build_v1_span() == SpanBuilder(
trace_id=trace_id,
name='GET /foo',
parent_id=parent_span_id,
span_id=server_span_id,
timestamp=24.0,
duration=18.0,
annotations={
'sr': 24,
'ss': 42
},
tags={
'k': 'v'
},
kind=Kind.SERVER,
service_name=client_svc_name,
local_endpoint=fake_endpoint,
).build_v1_span()
assert client_log_call[0][1] == client_span
assert flush_mock.call_count == 1
def test_get_thread_local_span_storage_present():
tracer = storage.get_default_tracer()
with mock.patch.object(tracer, "_span_storage", SpanStorage(["foo"])):
assert thread_local.get_thread_local_span_storage() == SpanStorage(
["foo"])
def test_span_context(
generate_string_128bit_mock,
generate_string_mock,
thread_local_mock,
span_func,
expected_annotations,
):
span_storage = SpanStorage()
generate_string_mock.return_value = '1'
with zipkin.zipkin_span(
service_name='root_span',
span_name='root_span',
sample_rate=100.0,
transport_handler=MockTransportHandler(),
span_storage=span_storage,
):
zipkin_attrs = ZipkinAttrs(
trace_id='1111111111111111',
span_id='2222222222222222',
parent_span_id='3333333333333333',
flags='flags',
is_sampled=True,
)
thread_local_mock.zipkin_attrs = [zipkin_attrs]
ts = time.time()
with mock.patch('time.time', return_value=ts):
with span_func(
service_name='svc',
span_name='span',
binary_annotations={'foo': 'bar'},
span_storage=span_storage,
):
# Assert that the new ZipkinAttrs were saved
new_zipkin_attrs = get_zipkin_attrs()
assert new_zipkin_attrs.span_id == '1'
# Outside of the context, things should be returned to normal
assert get_zipkin_attrs() == zipkin_attrs
client_span = span_storage.pop().build_v1_span()
# These reserved annotations are based on timestamps so pop em.
# This also acts as a check that they exist.
for annotation in expected_annotations:
client_span.annotations.pop(annotation)
expected_client_span = _V1Span(
trace_id='1111111111111111',
name='span',
parent_id='2222222222222222',
id='1',
timestamp=ts,
duration=0.0,
endpoint=None,
annotations={},
binary_annotations={'foo': 'bar'},
sa_endpoint=None,
)
assert client_span == expected_client_span
assert generate_string_128bit_mock.call_count == 0