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_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, 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",
)
with zipkin.zipkin_span(
service_name="test_service_name",
span_name="producer_span",
timestamp=ts,
duration=10,
kind=Kind.PRODUCER,
):
pass
output = mock_transport_handler.get_payloads()[0]
validate_fn(output, zipkin_attrs, inner_span_id, ts)
def _do_test_concurrent_subrequests_in_threads(thread_class):
"""What if a thread has a span context, then fires off N threads to do N
subrequests in parallel?
Surely the spans of the subrequests should be siblings whose parentID is
equal to the first thread's span context's spanID.
Furthermore, a final sub-span created in the first thread should be a child
of that thread (no leaking of span stack from the child threads).
"""
transport = MockTransportHandler()
with zipkin_span(
service_name="main_thread",
span_name="handle_one_big_request",
transport_handler=transport,
sample_rate=100.0,
encoding=Encoding.V2_JSON,
use_128bit_trace_id=True,
) as span_ctx:
assert True is span_ctx._is_local_root_span
assert span_ctx.logging_context
expected_trace_id = span_ctx.zipkin_attrs.trace_id
expected_parent_id = span_ctx.zipkin_attrs.span_id
# Now do three subrequests
req_count = 3
threads = []
output_q = queue.Queue()
for thread_idx in range(req_count):
this_thread = thread_class(
target=_do_one_little_request, args=(output_q, "input-%d" % thread_idx)
)
threads.append(this_thread)
this_thread.start()
outputs = set()
for thread in threads:
thread.join()
outputs.add(output_q.get())
assert {"input-0-output", "input-1-output", "input-2-output"} == outputs
output = transport.get_payloads()
assert len(output) == 1
spans = sorted(json.loads(output[0]), key=itemgetter("timestamp"))
assert len(spans) == 4
parent_span = spans[0]
subrequest_spans = spans[1:]
assert len(subrequest_spans) == 3
assert parent_span["name"] == "handle_one_big_request"
for span in subrequest_spans:
assert "do-the-thing" == span["name"]
assert span["tags"]["input.was"] in ("input-0", "input-1", "input-2")
assert expected_trace_id == span["traceId"]
# Perhaps most importantly, all the subrequest spans should share the same
# parentId, which is the main thread's span's 'id'
assert expected_parent_id == span["parentId"]
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__get_path_content_type(self, encoding, path, content_type):
transport = MockTransportHandler()
with zipkin_span(
service_name="my_service",
span_name="home",
sample_rate=100,
transport_handler=transport,
encoding=encoding,
):
pass
spans = transport.get_payloads()[0]
http_transport = SimpleHTTPTransport("localhost", 9411)
assert http_transport._get_path_content_type(spans) == (path,
content_type)
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_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_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_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_many_times(fake_endpoint):
# We create MAX_PORTION_SIZE * 2 + 1 spans, so we should trigger flush 3
# times, once every MAX_PORTION_SIZE spans.
encoder = MockEncoder()
sender = logging_helper.ZipkinBatchSender(
transport_handler=MockTransportHandler(),
max_portion_size=None,
encoder=encoder,
)
max_portion_size = logging_helper.ZipkinBatchSender.MAX_PORTION_SIZE
with sender:
for _ in range(max_portion_size * 2 + 1):
sender.add_span(
Span(
trace_id="000000000000000f",
name="span",
parent_id="0000000000000001",
span_id="0000000000000002",
kind=Kind.CLIENT,
timestamp=26.0,
duration=4.0,
local_endpoint=fake_endpoint,
annotations={},
tags={},
)
)
assert encoder.encode_queue.call_count == 3
assert len(encoder.encode_queue.call_args_list[0][0][0]) == max_portion_size
assert len(encoder.encode_queue.call_args_list[1][0][0]) == max_portion_size
assert len(encoder.encode_queue.call_args_list[2][0][0]) == 1
def test_get_current_context_root_sample_rate_override_not_sampled(
self,
mock_create_attr,
):
# Root span, with custom zipkin_attrs, not sampled and sample_rate
zipkin_attrs = ZipkinAttrs(
trace_id=generate_random_64bit_string(),
span_id=generate_random_64bit_string(),
parent_span_id=generate_random_64bit_string(),
flags=None,
is_sampled=False,
)
context = zipkin.zipkin_span(
service_name='test_service',
span_name='test_span',
transport_handler=MockTransportHandler(),
zipkin_attrs=zipkin_attrs,
sample_rate=100.0,
)
report_root, _ = context._get_current_context()
assert mock_create_attr.call_args == mock.call(
sample_rate=100.0,
trace_id=zipkin_attrs.trace_id,
)
# It wasn't sampled before and now it is, so this is the trace root
assert report_root is True
def test_get_current_context_root_sample_rate_override_sampled(
self,
mock_create_attr,
):
# Root span, with custom zipkin_attrs, sampled
# Just return the custom zipkin_attrs.
zipkin_attrs = ZipkinAttrs(
trace_id=generate_random_64bit_string(),
span_id=generate_random_64bit_string(),
parent_span_id=generate_random_64bit_string(),
flags=None,
is_sampled=True,
)
context = zipkin.zipkin_span(
service_name='test_service',
span_name='test_span',
transport_handler=MockTransportHandler(),
zipkin_attrs=zipkin_attrs,
sample_rate=100.0,
)
report_root, current_attrs = context._get_current_context()
assert mock_create_attr.call_count == 0
assert current_attrs == zipkin_attrs
# The override was set and was already sampled, so this is probably
# not the trace root.
assert report_root is False
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_batch_sender_with_error_on_exit():
sender = logging_helper.ZipkinBatchSender(
MockTransportHandler(), None, MockEncoder(),
)
with pytest.raises(ZipkinError):
with sender:
raise Exception("Error!")
def test_batch_sender_add_span(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(
Span(
trace_id="000000000000000f",
name="span",
parent_id="0000000000000001",
span_id="0000000000000002",
kind=Kind.CLIENT,
timestamp=26.0,
duration=4.0,
local_endpoint=fake_endpoint,
annotations={},
tags={},
)
)
assert encoder.encode_queue.call_count == 1
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_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_decorator_works_in_a_new_thread():
"""The zipkin_span decorator is instantiated in a thread and then run in
another. Let's verify that it works and that it stores the span in the
right thread's thread-storage.
"""
transport = MockTransportHandler()
thread = threading.Thread(target=run_inside_another_thread, args=(transport,))
thread.start()
thread.join()
output = transport.get_payloads()
assert len(output) == 1
spans = json.loads(output[0])
assert len(spans) == 2
assert spans[0]["name"] == "service1_do_stuff"
assert spans[1]["name"] == "index"
def test_initinvalid_sample_rate(self):
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_init_local_root_transport_zipkin_attrs(self):
# If transport_handler and zipkin_attrs are set, then this is a local root
context = zipkin.zipkin_span(
service_name='test_service',
span_name='test_span',
transport_handler=MockTransportHandler(),
zipkin_attrs=zipkin.create_attrs_for_span(),
)
assert context._is_local_root_span is True
def test_init_local_root_transport_sample_rate(self):
# If transport_handler and sample_rate are set, then this is a local root
context = zipkin.zipkin_span(
service_name='test_service',
span_name='test_span',
transport_handler=MockTransportHandler(),
sample_rate=100.0,
)
assert context._is_local_root_span is True
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_init_local_root_firehose(self):
# If firehose_handler is set, then this is a local root even if there's
# no sample_rate or zipkin_attrs
context = zipkin.zipkin_span(
service_name='test_service',
span_name='test_span',
firehose_handler=MockTransportHandler(),
)
assert context._is_local_root_span is True
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_override_span_name(self):
with zipkin.zipkin_client_span(
service_name='test_service',
span_name='test_span',
transport_handler=MockTransportHandler(),
sample_rate=100.0,
) as span:
span.override_span_name('new_name')
assert span.span_name == 'new_name'
assert span.logging_context.span_name == 'new_name'
def test_start_root_span_not_sampled_firehose(self, mock_log_ctx):
# This request is not sampled, but firehose is setup. So we need to
# setup the transport anyway.
transport = MockTransportHandler()
firehose = MockTransportHandler()
tracer = MockTracer()
context = tracer.zipkin_span(
service_name='test_service',
span_name='test_span',
transport_handler=transport,
firehose_handler=firehose,
sample_rate=0.0,
)
context.start()
assert context.zipkin_attrs is not None
assert mock_log_ctx.call_count == 1
assert mock_log_ctx.return_value.start.call_count == 1
assert tracer.is_transport_configured() is True
def test_override_span_name(self):
with zipkin.zipkin_client_span(
service_name="test_service",
span_name="test_span",
transport_handler=MockTransportHandler(),
sample_rate=100.0,
) as span:
span.override_span_name("new_name")
assert span.span_name == "new_name"
assert span.logging_context.span_name == "new_name"
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_init_span_storage_wrong_type(self):
# 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 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'