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()