def test_zipkin_client_span():
context = zipkin.zipkin_client_span('test_service', 'test_span')
assert context.kind == Kind.CLIENT
with pytest.raises(ValueError):
zipkin.zipkin_client_span('test_service', 'test_span', kind=Kind.LOCAL)
def test_can_set_sa_annotation(encoding):
mock_transport_handler, mock_logs = mock_logger()
with zipkin.zipkin_client_span(
service_name='test_service_name',
span_name='test_span_name',
transport_handler=mock_transport_handler,
sample_rate=100.0,
encoding=encoding,
) as span:
span.add_sa_binary_annotation(
port=8888,
service_name='sa_service',
host='10.0.0.0',
)
assert len(mock_logs) == 1
client_span = json.loads(mock_logs[0])[0]
if encoding == Encoding.V1_JSON:
assert client_span['binaryAnnotations'][0]['key'] == 'sa'
assert client_span['binaryAnnotations'][0]['value'] is True
host = client_span['binaryAnnotations'][0]['endpoint']
elif encoding == Encoding.V2_JSON:
host = client_span['remoteEndpoint']
assert host['serviceName'] == u'sa_service'
assert host['ipv4'] == '10.0.0.0'
assert 'ipv6' not in host
assert host['port'] == 8888
def test_can_set_sa_annotation(encoding):
mock_transport_handler, mock_logs = mock_logger()
with zipkin.zipkin_client_span(
service_name="test_service_name",
span_name="test_span_name",
transport_handler=mock_transport_handler,
sample_rate=100.0,
encoding=encoding,
) as span:
span.add_sa_binary_annotation(
port=8888,
service_name="sa_service",
host="10.0.0.0",
)
assert len(mock_logs) == 1
client_span = json.loads(mock_logs[0])[0]
if encoding == Encoding.V1_JSON:
assert client_span["binaryAnnotations"][0]["key"] == "sa"
assert client_span["binaryAnnotations"][0]["value"] is True
host = client_span["binaryAnnotations"][0]["endpoint"]
elif encoding == Encoding.V2_JSON:
host = client_span["remoteEndpoint"]
assert host["serviceName"] == u"sa_service"
assert host["ipv4"] == "10.0.0.0"
assert "ipv6" not in host
assert host["port"] == 8888
def execute(self, sql, params=()):
with zipkin_client_span(
service_name=self.get_db_type(),
span_name=self.get_operation(sql),
binary_annotations=self.get_binary_annotations(sql),
) as span:
self.add_sa_binary_annotation(span)
return self.cursor.execute(sql, params)
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_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_add_sa_binary_annotation_non_root(self):
# Nothing happens if this is not a client span
with zipkin.zipkin_client_span('test_service', 'test_span') as span:
span.add_sa_binary_annotation(80, 'remote_service', '10.1.2.3')
expected_endpoint = create_endpoint(80, 'remote_service', '10.1.2.3')
assert span.remote_endpoint == expected_endpoint
# setting it again throws an error
with pytest.raises(ValueError):
span.add_sa_binary_annotation(80, 'remote_service', '10.1.2.3')
def _call_traced(self: MethodProxy, *args, **kwargs):
span = zipkin.zipkin_client_span(self.service_name,
self.method_name,
transport_handler=transport_handler)
start_span(span)
self.worker_ctx.data.update(zipkin.create_http_headers_for_new_span())
try:
reply = _call(self, *args, **kwargs)
except:
stop_span(span)
raise
return TracedRpcReply(reply.reply_event, span)
def test_server_span(encoding):
mock_transport_handler, mock_logs = mock_logger()
with zipkin.zipkin_server_span(
service_name="test_service_name",
span_name="test_span_name",
transport_handler=mock_transport_handler,
sample_rate=100.0,
encoding=encoding,
):
with zipkin.zipkin_client_span(
service_name="nested_service",
span_name="nested_span",
):
pass
pass
assert len(mock_logs) == 1
spans = json.loads(mock_logs[0])
client_span = spans[0]
server_span = spans[1]
assert server_span["name"] == "test_span_name"
# Local nested spans report timestamp and duration
assert "timestamp" in server_span
assert "duration" in server_span
if encoding == Encoding.V1_JSON:
assert len(server_span["annotations"]) == 2
assert sorted([ann["value"]
for ann in server_span["annotations"]]) == [
"sr",
"ss",
]
elif encoding == Encoding.V2_JSON:
assert server_span["kind"] == "SERVER"
assert client_span["name"] == "nested_span"
# Local nested spans report timestamp and duration
assert "timestamp" in client_span
assert "duration" in client_span
if encoding == Encoding.V1_JSON:
assert len(client_span["annotations"]) == 2
assert sorted([ann["value"]
for ann in client_span["annotations"]]) == [
"cr",
"cs",
]
elif encoding == Encoding.V2_JSON:
assert client_span["kind"] == "CLIENT"
def test_add_sa_binary_annotation_root(self):
# Nothing happens if this is not a client span
with zipkin.zipkin_client_span(
service_name='test_service',
span_name='test_span',
transport_handler=MockTransportHandler(),
sample_rate=100.0,
) as span:
span.add_sa_binary_annotation(80, 'remote_service', '10.1.2.3')
expected_endpoint = create_endpoint(80, 'remote_service', '10.1.2.3')
assert span.logging_context.remote_endpoint == expected_endpoint
# setting it again throws an error
with pytest.raises(ValueError):
span.add_sa_binary_annotation(80, 'remote_service', '10.1.2.3')
def test_server_span(encoding):
mock_transport_handler, mock_logs = mock_logger()
with zipkin.zipkin_server_span(
service_name='test_service_name',
span_name='test_span_name',
transport_handler=mock_transport_handler,
sample_rate=100.0,
encoding=encoding,
):
with zipkin.zipkin_client_span(
service_name='nested_service',
span_name='nested_span',
):
pass
pass
assert len(mock_logs) == 1
spans = json.loads(mock_logs[0])
client_span = spans[0]
server_span = spans[1]
assert server_span['name'] == 'test_span_name'
# Local nested spans report timestamp and duration
assert 'timestamp' in server_span
assert 'duration' in server_span
if encoding == Encoding.V1_JSON:
assert len(server_span['annotations']) == 2
assert sorted([ann['value']
for ann in server_span['annotations']]) == ['sr', 'ss']
elif encoding == Encoding.V2_JSON:
assert server_span['kind'] == 'SERVER'
assert client_span['name'] == 'nested_span'
# Local nested spans report timestamp and duration
assert 'timestamp' in client_span
assert 'duration' in client_span
if encoding == Encoding.V1_JSON:
assert len(client_span['annotations']) == 2
assert sorted([ann['value']
for ann in client_span['annotations']]) == ['cr', 'cs']
elif encoding == Encoding.V2_JSON:
assert client_span['kind'] == 'CLIENT'
def _start_rpc_span(self, sender, zipkin_attrs):
# type: (str, py_zipkin.zipkin.ZipkinAttrs) -> py_zipkin.zipkin.client_span
zipkin_attrs = zipkin.ZipkinAttrs(
trace_id=zipkin_attrs.trace_id,
span_id=zipkin.generate_random_64bit_string(),
parent_span_id=zipkin_attrs.span_id,
flags=zipkin_attrs.flags,
is_sampled=zipkin_attrs.is_sampled,
)
span = zipkin.zipkin_client_span(
service_name='celery.{}'.format(sender.split('.', 1)[0]),
span_name='task.rpc.{}'.format(sender),
transport_handler=self.transport_handler,
sample_rate=self.sample_rate,
zipkin_attrs=zipkin_attrs,
)
span.start()
return span
def test_memory_leak():
# In py_zipkin >= 0.13.0 and <= 0.14.0 this test fails since the
# span_storage contains 10 spans once you exit the for loop.
mock_transport_handler, mock_logs = mock_logger()
assert len(get_default_tracer().get_spans()) == 0
for _ in range(10):
with zipkin.zipkin_client_span(
service_name='test_service_name',
span_name='test_span_name',
transport_handler=mock_transport_handler,
sample_rate=0.0,
encoding=Encoding.V2_JSON,
):
with zipkin.zipkin_span(
service_name='inner_service_name',
span_name='inner_span_name',
):
pass
assert len(mock_logs) == 0
assert len(get_default_tracer().get_spans()) == 0
def test_client_span(encoding):
"""Tests the zipkin_client_span helper."""
mock_transport_handler, mock_logs = mock_logger()
with zipkin.zipkin_client_span(
service_name='test_service_name',
span_name='test_span_name',
transport_handler=mock_transport_handler,
sample_rate=100.0,
encoding=encoding,
):
pass
assert len(mock_logs) == 1
span = json.loads(mock_logs[0])[0]
assert span['name'] == 'test_span_name'
if encoding == Encoding.V1_JSON:
assert sorted([ann['value'] for ann in span['annotations']]) == \
['cr', 'cs']
elif encoding == Encoding.V2_JSON:
assert span['kind'] == 'CLIENT'
def test_client_span(encoding):
"""Tests the zipkin_client_span helper."""
mock_transport_handler, mock_logs = mock_logger()
with zipkin.zipkin_client_span(
service_name="test_service_name",
span_name="test_span_name",
transport_handler=mock_transport_handler,
sample_rate=100.0,
encoding=encoding,
):
pass
assert len(mock_logs) == 1
span = json.loads(mock_logs[0])[0]
assert span["name"] == "test_span_name"
if encoding == Encoding.V1_JSON:
assert sorted([ann["value"]
for ann in span["annotations"]]) == ["cr", "cs"]
elif encoding == Encoding.V2_JSON:
assert span["kind"] == "CLIENT"
def urlopen(self, method, url, **kw):
if parsed_host and self.host == parsed_host.hostname:
# Don't trace zipkin calls
return func(self, method, url, **kw)
attrs = {
zipkin_core.HTTP_HOST: self.host,
zipkin_core.HTTP_METHOD: method,
zipkin_core.HTTP_PATH: url,
}
with zipkin_client_span(
service_name=self.host, span_name=self._absolute_url(url), binary_annotations=attrs
) as span:
headers = kw.pop("headers", {})
headers.update(create_http_headers_for_new_span())
try:
out = func(self, method, url, headers=headers, **kw)
if hasattr(out.connection, "sock") and hasattr(out.connection.sock, "getpeername"):
peer = out.connection.sock.getpeername()
span.add_sa_binary_annotation(peer[1], self.host, peer[0])
else:
span.add_sa_binary_annotation(self.port, self.host)
except:
# always add sa_binary even in case of error
# but if we do it before firing urlopen, then we ended up with two annotations
span.add_sa_binary_annotation(self.port, self.host)
raise
span.update_binary_annotations(
{zipkin_core.HTTP_STATUS_CODE: out.status, "http.retries": out.retries.total,}
)
return out
