def test_encoding(encoding, validate_fn):
zipkin_attrs = ZipkinAttrs(
trace_id=generate_random_64bit_string(),
span_id=generate_random_64bit_string(),
parent_span_id=generate_random_64bit_string(),
is_sampled=True,
flags=None,
)
inner_span_id = generate_random_64bit_string()
mock_transport_handler = MockTransportHandler(10000)
# Let's hardcode the timestamp rather than call time.time() every time.
# The issue with time.time() is that the convertion to int of the
# returned float value * 1000000 is not precise and in the same test
# sometimes returns N and sometimes N+1. This ts value doesn't have that
# issue afaict, probably since it ends in zeros.
ts = 1538544126.115900
with mock.patch('time.time', autospec=True) as mock_time:
# zipkin.py start, logging_helper.start, 3 x logging_helper.stop
# I don't understand why logging_helper.stop would run 3 times, but
# that's what I'm seeing in the test
mock_time.side_effect = iter([ts, ts, ts + 10, ts + 10, ts + 10])
with zipkin.zipkin_span(
service_name='test_service_name',
span_name='test_span_name',
transport_handler=mock_transport_handler,
binary_annotations={'some_key': 'some_value'},
encoding=encoding,
zipkin_attrs=zipkin_attrs,
host='10.0.0.0',
port=8080,
kind=Kind.CLIENT,
) as span:
with mock.patch.object(
zipkin,
'generate_random_64bit_string',
return_value=inner_span_id,
):
with zipkin.zipkin_span(
service_name='test_service_name',
span_name='inner_span',
timestamp=ts,
duration=5,
annotations={'ws': ts},
):
span.add_sa_binary_annotation(
8888,
'sa_service',
'2001:0db8:85a3:0000:0000:8a2e:0370:7334',
)
output = mock_transport_handler.get_payloads()[0]
validate_fn(output, zipkin_attrs, inner_span_id, ts)
def _create_root_span(is_sampled, firehose_enabled):
return (), {
'root_span':
zipkin.zipkin_span(
service_name='my_service',
span_name='my_span_name',
transport_handler=MockTransportHandler(),
firehose_handler=MockTransportHandler()
if firehose_enabled else None,
port=42,
sample_rate=0 if not is_sampled else 100,
)
}
def test_batch_sender_add_span_many_times(
mock_encode_bytes_list,
empty_annotations_dict,
empty_binary_annotations_dict,
fake_endpoint,
):
# We create MAX_PORTION_SIZE * 2 + 1 spans, so we should trigger flush 3
# times, once every MAX_PORTION_SIZE spans.
sender = logging_helper.ZipkinBatchSender(MockTransportHandler())
max_portion_size = logging_helper.ZipkinBatchSender.MAX_PORTION_SIZE
with sender:
for _ in range(max_portion_size * 2 + 1):
sender.add_span(
span_id='0000000000000002',
parent_span_id='0000000000000001',
trace_id='000000000000000f',
span_name='span',
annotations=empty_annotations_dict,
binary_annotations=empty_binary_annotations_dict,
timestamp_s=None,
duration_s=None,
endpoint=fake_endpoint,
sa_endpoint=None,
)
assert mock_encode_bytes_list.call_count == 3
assert len(
mock_encode_bytes_list.call_args_list[0][0][0]) == max_portion_size
assert len(
mock_encode_bytes_list.call_args_list[1][0][0]) == max_portion_size
assert len(mock_encode_bytes_list.call_args_list[2][0][0]) == 1
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(
empty_annotations_dict,
empty_binary_annotations_dict,
fake_endpoint,
):
# This test verifies it's possible to add 1 span without throwing errors.
# It also checks that exiting the ZipkinBatchSender context manager
# triggers a flush of all the already added spans.
encoder = MockEncoder(encoded_queue='foobar')
sender = logging_helper.ZipkinBatchSender(
transport_handler=MockTransportHandler(),
max_portion_size=None,
encoder=encoder,
)
with sender:
sender.add_span(
SpanBuilder(
trace_id='000000000000000f',
name='span',
parent_id='0000000000000001',
span_id='0000000000000002',
timestamp=26.0,
duration=4.0,
annotations=empty_annotations_dict,
tags=empty_binary_annotations_dict,
kind=Kind.CLIENT,
local_endpoint=fake_endpoint,
service_name='test_service',
))
assert encoder.encode_queue.call_count == 1
def test_span_context_sampled_no_handlers(
generate_string_mock,
thread_local_mock,
):
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]
generate_string_mock.return_value = '1'
context = zipkin.zipkin_span(
service_name='my_service',
port=5,
transport_handler=MockTransportHandler(),
sample_rate=None,
)
with context:
# Assert that the new ZipkinAttrs were saved
new_zipkin_attrs = ThreadLocalStack().get()
assert new_zipkin_attrs.span_id == '1'
# Outside of the context, things should be returned to normal
assert ThreadLocalStack().get() == zipkin_attrs
def test_call_calls_send(self):
with mock.patch.object(MockTransportHandler, 'send', autospec=True):
transport = MockTransportHandler()
transport("foobar")
assert transport.send.call_count == 1
assert transport.send.call_args == mock.call(transport, "foobar")
def test_update_binary_annotations():
zipkin_attrs = ZipkinAttrs(
trace_id='0',
span_id='1',
parent_span_id=None,
flags='0',
is_sampled=True,
)
context = zipkin.zipkin_span(
service_name='my_service',
span_name='span_name',
zipkin_attrs=zipkin_attrs,
transport_handler=MockTransportHandler(),
port=5,
)
with context:
assert 'test' not in context.logging_context.binary_annotations_dict
context.update_binary_annotations({'test': 'hi'})
assert context.logging_context.binary_annotations_dict['test'] == 'hi'
nested_context = zipkin.zipkin_span(
service_name='my_service',
span_name='nested_span',
binary_annotations={'one': 'one'},
)
with nested_context:
assert 'one' not in context.logging_context.binary_annotations_dict
nested_context.update_binary_annotations({'two': 'two'})
assert 'two' in nested_context.binary_annotations
assert 'two' not in context.logging_context.binary_annotations_dict
def test_adding_error_annotation_on_exception(
mock_update_binary_annotations,
exception_message,
expected_error_string,
):
zipkin_attrs = ZipkinAttrs(
trace_id='0',
span_id='1',
parent_span_id=None,
flags='0',
is_sampled=True,
)
context = zipkin.zipkin_span(
service_name='my_service',
span_name='span_name',
zipkin_attrs=zipkin_attrs,
transport_handler=MockTransportHandler(),
port=5,
)
with pytest.raises(ValueError):
with context:
raise ValueError(exception_message)
assert mock_update_binary_annotations.call_count == 1
call_args, _ = mock_update_binary_annotations.call_args
assert 'error' in call_args[1]
assert expected_error_string == call_args[1]['error']
def test_zipkin_invalid_include():
with pytest.raises(ZipkinError):
with zipkin.zipkin_span(service_name='some_service_name',
span_name='span_name',
transport_handler=MockTransportHandler(),
sample_rate=100.0,
include=('clawyant', )):
pass
def test_zipkin_extraneous_include_raises(mock_zipkin_span, span_func):
with pytest.raises(ValueError):
with span_func(service_name='some_service_name',
span_name='span_name',
transport_handler=MockTransportHandler(),
sample_rate=100.0,
include=('foobar', )):
assert mock_zipkin_span.__init__.call_count == 0
def test_zipkin_invalid_sample_rate():
with pytest.raises(ZipkinError):
with zipkin.zipkin_span(
service_name='some_service_name',
span_name='span_name',
transport_handler=MockTransportHandler(),
sample_rate=101.0,
):
pass
with pytest.raises(ZipkinError):
with zipkin.zipkin_span(
service_name='some_service_name',
span_name='span_name',
transport_handler=MockTransportHandler(),
sample_rate=-0.1,
):
pass
def test_batch_sender_with_error_on_exit():
sender = logging_helper.ZipkinBatchSender(
MockTransportHandler(),
None,
MockEncoder(),
)
with pytest.raises(ZipkinError):
with sender:
raise Exception('Error!')
def test_zipkin_span_add_logging_annotation(mock_context):
with zipkin.zipkin_span(
service_name='my_service',
transport_handler=MockTransportHandler(),
sample_rate=100.0,
add_logging_annotation=True,
):
pass
_, kwargs = mock_context.call_args
assert kwargs['add_logging_annotation']
def test_zipkin_extraneous_kind_raises(mock_zipkin_span, span_func):
with pytest.raises(ValueError):
with span_func(
service_name='some_service_name',
span_name='span_name',
transport_handler=MockTransportHandler(),
sample_rate=100.0,
kind=Kind.LOCAL,
):
pass
def test_zipkin_span_span_storage_wrong_type():
# Missing transport_handler
with pytest.raises(ZipkinError):
with zipkin.zipkin_span(
service_name='some_service_name',
span_name='span_name',
transport_handler=MockTransportHandler(),
port=5,
sample_rate=100.0,
span_storage=[],
):
pass
def context():
attr = ZipkinAttrs(None, None, None, None, False)
log_handler = logging_helper.ZipkinLoggerHandler(attr)
return logging_helper.ZipkinLoggingContext(
zipkin_attrs=attr,
endpoint=_encoding_helpers.create_endpoint(80, 'test_server',
'127.0.0.1'),
log_handler=log_handler,
span_name='span_name',
transport_handler=MockTransportHandler(),
report_root_timestamp=False,
)
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_override_span_name():
transport = MockTransportHandler()
with zipkin.zipkin_span(
service_name='my_service',
span_name='span_name',
transport_handler=transport,
kind=Kind.CLIENT,
sample_rate=100.0,
encoding=Encoding.V1_JSON,
) as context:
context.override_span_name('new_span_name')
with zipkin.zipkin_span(
service_name='my_service',
span_name='nested_span',
) as nested_context:
nested_context.override_span_name('new_nested_span')
spans = json.loads(transport.get_payloads()[0])
assert len(spans) == 2
assert spans[0]['name'] == 'new_nested_span'
assert spans[1]['name'] == 'new_span_name'
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_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_add_sa_binary_annotation():
zipkin_attrs = ZipkinAttrs(
trace_id='0',
span_id='1',
parent_span_id=None,
flags='0',
is_sampled=True,
)
context = zipkin.zipkin_span(
service_name='my_service',
span_name='span_name',
zipkin_attrs=zipkin_attrs,
transport_handler=MockTransportHandler(),
port=5,
kind=Kind.CLIENT,
)
with context:
assert context.logging_context.sa_endpoint is None
context.add_sa_binary_annotation(
port=123,
service_name='test_service',
host='1.2.3.4',
)
expected_sa_endpoint = create_endpoint(
port=123,
service_name='test_service',
host='1.2.3.4',
)
assert context.logging_context.sa_endpoint == \
expected_sa_endpoint
nested_context = zipkin.zipkin_span(
service_name='my_service',
span_name='nested_span',
kind=Kind.CLIENT,
)
with nested_context:
nested_context.add_sa_binary_annotation(
port=456,
service_name='nested_service',
host='5.6.7.8',
)
expected_nested_sa_endpoint = create_endpoint(
port=456,
service_name='nested_service',
host='5.6.7.8',
)
assert nested_context.sa_endpoint == \
expected_nested_sa_endpoint
def test_adding_sa_binary_annotation_for_non_client_spans():
context = zipkin.zipkin_span(
service_name='my_service',
span_name='span_name',
transport_handler=MockTransportHandler(),
include=('server', ),
sample_rate=100.0,
)
with context:
context.add_sa_binary_annotation(
port=123,
service_name='test_service',
host='1.2.3.4',
)
assert context.logging_context.sa_endpoint is None
def test_zipkin_invalid_kind():
with pytest.raises(ZipkinError):
with zipkin.zipkin_span(
service_name='some_service_name',
span_name='span_name',
transport_handler=MockTransportHandler(),
sample_rate=100.0,
):
with zipkin.zipkin_span(
service_name='nested_service',
span_name='nested_span',
kind='client',
):
pass
pass
def test_zipkin_span_decorator_many(create_endpoint_mock):
@zipkin.zipkin_span(service_name='decorator')
def test_func(a, b):
return a + b
assert test_func(1, 2) == 3
assert create_endpoint_mock.call_count == 0
with zipkin.zipkin_span(
service_name='context_manager',
transport_handler=MockTransportHandler(),
sample_rate=100.0,
):
assert test_func(1, 2) == 3
assert create_endpoint_mock.call_count == 1
assert test_func(1, 2) == 3
assert create_endpoint_mock.call_count == 1
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_add_sa_binary_annotation_twice():
zipkin_attrs = ZipkinAttrs(
trace_id='0',
span_id='1',
parent_span_id=None,
flags='0',
is_sampled=True,
)
context = zipkin.zipkin_span(
service_name='my_service',
span_name='span_name',
zipkin_attrs=zipkin_attrs,
transport_handler=MockTransportHandler(),
port=5,
)
with context:
assert context.logging_context.sa_endpoint is None
context.add_sa_binary_annotation(
port=123,
service_name='test_service',
host='1.2.3.4',
)
with pytest.raises(ValueError):
context.add_sa_binary_annotation(
port=123,
service_name='test_service',
host='1.2.3.4',
)
nested_context = zipkin.zipkin_span(
service_name='my_service',
span_name='nested_span',
)
with nested_context:
nested_context.add_sa_binary_annotation(
port=456,
service_name='nested_service',
host='5.6.7.8',
)
with pytest.raises(ValueError):
nested_context.add_sa_binary_annotation(
port=456,
service_name='nested_service',
host='5.6.7.8',
)
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_update_binary_annotations_should_not_error_if_not_tracing():
zipkin_attrs = ZipkinAttrs(
trace_id='0',
span_id='1',
parent_span_id=None,
flags='0',
is_sampled=False,
)
context = zipkin.zipkin_span(
service_name='my_service',
span_name='span_name',
zipkin_attrs=zipkin_attrs,
transport_handler=MockTransportHandler(),
port=5,
)
with context:
# A non-sampled request should result in a no-op
context.update_binary_annotations({'test': 'hi'})
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()