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