def test_sync_no_inject(app, tracer, b3_keys):
config = ZipkinConfig(inject_response_headers=False)
app.add_middleware(ZipkinMiddleware, config=config, _tracer=tracer)
client = TestClient(app)
response = client.get("/sync-message?foo=bar")
assert response.status_code == 200
assert not any(key in response.headers for key in b3_keys)
def test_sync(app, tracer, b3_keys):
config = ZipkinConfig(header_formatter=Headers)
app.add_middleware(ZipkinMiddleware, config=config, _tracer=tracer)
client = TestClient(app)
response = client.get("/sync-message?foo=bar")
assert response.status_code == 200
assert all(key in response.headers for key in b3_keys)
def test_split_char(app, tracer, uber_keys, split_char):
config = ZipkinConfig(header_formatter=Headers,
header_formatter_kwargs=dict(split_char=split_char))
app.add_middleware(ZipkinMiddleware, config=config, _tracer=tracer)
client = TestClient(app)
response = client.get("/sync-message?foo=bar")
assert response.status_code == 200
assert all(key in response.headers for key in uber_keys)
item = response.headers.get("uber-trace-id")
assert split_char in item
async def test_dispatch_trace_reuse_tracer(app, dummy_request, next_response):
config = ZipkinConfig()
middleware = ZipkinMiddleware(app, config=config)
# the tracer is initialized on the first dispatch
assert middleware.tracer is None
await middleware.dispatch(dummy_request(), next_response)
assert middleware.tracer is not None
tracer = middleware.tracer
await middleware.dispatch(dummy_request(), next_response)
assert middleware.tracer is tracer, "Tracer must be reused on every requests"
await tracer.close()
def test_sync_force_new_trace(app, tracer, b3_keys):
config = ZipkinConfig(force_new_trace=True)
app.add_middleware(ZipkinMiddleware, config=config, _tracer=tracer)
client = TestClient(app)
response = client.get("/sync-message?foo=bar")
# call with injected tracing headers - needs to follow up
headers = {}
for key in b3_keys:
headers[key] = response.headers[key]
response2 = client.get("/sync-message?foo=bar", headers=headers)
assert response2.status_code == 200
assert "x-b3-parentspanid" not in response2.headers
assert headers["x-b3-traceid"] != response2.headers["x-b3-traceid"]
def test_async_request_data(app, tracer, b3_keys):
config = ZipkinConfig(header_formatter=Headers)
app.add_middleware(ZipkinMiddleware, config=config, _tracer=tracer)
client = TestClient(app)
response = client.get("/async-message?foo=bar")
assert response.status_code == 200
assert all(key in response.headers for key in b3_keys)
# call with injected tracing headers - needs to follow up
headers = {}
for key in b3_keys:
headers[key] = response.headers[key]
response2 = client.get("/async-message?foo=bar", headers=headers)
assert response2.status_code == 200
assert all(key in response2.headers for key in b3_keys)
assert "x-b3-parentspanid" in response2.headers
assert (headers[Headers.TRACE_ID_HEADER] == response2.headers[
Headers.TRACE_ID_HEADER])
assert headers["x-b3-spanid"] == response2.headers["x-b3-parentspanid"]
async def test_dispatch_trace(app, dummy_request, next_response):
config = ZipkinConfig()
middleware = ZipkinMiddleware(app, config=config)
# the tracer is initialized on the first dispatch
assert middleware.tracer is None
resp = await middleware.dispatch(
dummy_request(headers={}),
next_response,
)
assert middleware.tracer is not None
assert middleware.tracer._transport is not None
assert (str(middleware.tracer._transport._address) ==
"http://localhost:9411/api/v2/spans")
assert dict(resp.headers) == {
"x-b3-flags": "0",
"x-b3-sampled": "1",
"x-b3-spanid": resp.headers["x-b3-spanid"],
"x-b3-traceid": resp.headers["x-b3-traceid"],
}
await middleware.tracer.close()
def test_async_request_data(app, tracer, uber_keys):
config = ZipkinConfig(header_formatter=Headers)
app.add_middleware(ZipkinMiddleware, config=config, _tracer=tracer)
client = TestClient(app)
response = client.get("/async-message?foo=bar")
assert response.status_code == 200
assert all(key in response.headers for key in uber_keys)
trace_id, span_id, parent_id, debug, sampled = Headers(
)._parse_uber_headers(response.headers)
# call with injected tracing headers - needs to follow up
headers = {}
for key in uber_keys:
headers[key] = response.headers[key]
response2 = client.get("/async-message?foo=bar", headers=headers)
assert response2.status_code == 200
assert all(key in response2.headers for key in uber_keys)
trace_id2, span_id2, parent_id2, debug2, sampled2 = Headers(
)._parse_uber_headers(response2.headers)
assert trace_id == trace_id2
assert span_id == parent_id2
async def test_dispatch_trace_new_child(app, dummy_request, next_response):
trace_id = "6223635aa7bfb6597d72ac7c4680bfed"
span_id = "ac7cb16943218de4"
config = ZipkinConfig("zipkin.host")
middleware = ZipkinMiddleware(app, config=config)
# the tracer is initialized on the first dispatch
assert middleware.tracer is None
resp = await middleware.dispatch(
dummy_request(headers={
"x-b3-spanid": span_id,
"x-b3-traceid": trace_id,
}),
next_response,
)
assert middleware.tracer is not None
assert middleware.tracer._transport is not None
assert (str(middleware.tracer._transport._address) ==
"http://zipkin.host:9411/api/v2/spans")
assert dict(resp.headers) == {
"x-b3-spanid": span_id,
"x-b3-traceid": trace_id,
}
await middleware.tracer.close()
async def test_dispatch_trace_buggy_headers(app, dummy_request, next_response):
trace_id = "6223635aa7bfb6597d72ac7c4680bfed"
config = ZipkinConfig()
middleware = ZipkinMiddleware(app, config=config)
# the tracer is initialized on the first dispatch
assert middleware.tracer is None
resp = await middleware.dispatch(
dummy_request(headers={
"x-b3-traceid": trace_id,
# x-b3-spanid should be here
}),
next_response,
)
assert middleware.tracer is not None
assert middleware.tracer._transport is not None
assert dict(resp.headers) == {
"x-b3-flags": "0",
"x-b3-sampled": "1",
"x-b3-spanid": resp.headers["x-b3-spanid"],
"x-b3-traceid": resp.headers["x-b3-traceid"],
}
# we cannot reuse the traceid if the span id was missing
assert trace_id != resp.headers["x-b3-spanid"]
await middleware.tracer.close()
# root span from middleware injects headers
# and becomes the parent for subsequet services
headers = child_span.context.make_headers()
child_span.name("NewParent")
child_span.annotate("Child, sleeps for 1, injects headers and becomes parent")
await asyncio.sleep(1)
return JSONResponse({"hello": "world"}, headers=headers)
routes = [
Route("/", JSONResponse({"status": "OK"})),
Route("/homepage", homepage),
]
app = Starlette(debug=True, routes=routes)
config = ZipkinConfig(
host="localhost",
port=9411,
service_name="service_name",
sample_rate=1.0,
inject_response_headers=True,
force_new_trace=False,
json_encoder=json.dumps,
header_formatter=UberHeaders,
)
app.add_middleware(ZipkinMiddleware, config=config)
if __name__ == "__main__":
uvicorn.run("app:app", host="0.0.0.0", port=8000, log_level="info", reload=True, workers=1)
@trace("api call", "CLIENT")
async def api_call():
async with AsyncClient() as cli:
return await cli.get("http://api:8000/", headers=trace.make_headers())
async def homepage(request):
with trace("before api"):
await asyncio.sleep(random())
response = await api_call()
with trace("after api", "PRODUCER"):
await asyncio.sleep(random())
return JSONResponse(response.json())
routes = [
Route("/", homepage),
]
app = Starlette(debug=True, routes=routes)
config = ZipkinConfig(host=TRACER, service_name="app")
app.add_middleware(ZipkinMiddleware, config=config)
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")
def test_get_transaction(app, params):
config = ZipkinConfig()
middleware = ZipkinMiddleware(app, config=config)
transac = middleware.get_transaction({"endpoint": params["endpoint"]})
assert transac == params["expected"]
def test_get_query(app, params):
config = ZipkinConfig()
middleware = ZipkinMiddleware(app, config=config)
querystring = middleware.get_query({"query_string": params["querystring"]})
assert querystring == params["expected"]
def test_get_headers(app, params):
config = ZipkinConfig()
middleware = ZipkinMiddleware(app, config=config)
headers = middleware.get_headers({"headers": params["headers"]})
assert headers == '{"a": "A, B"}'