Files
fastapi/tests/test_telemetry/test_background.py
T

277 lines
9.4 KiB
Python

import threading
from functools import partial
import anyio
import pytest
from fastapi import BackgroundTasks, Depends, FastAPI, HTTPException
from fastapi.responses import JSONResponse
from fastapi.telemetry import get_telemetry_data
from fastapi.testclient import TestClient
from opentelemetry import trace
from opentelemetry.sdk.trace import Span
from opentelemetry.sdk.trace.sampling import ALWAYS_OFF
from opentelemetry.trace import SpanKind, StatusCode
from starlette.background import BackgroundTask
from ._subprocess import run_in_subprocess
from .conftest import metric_points, server_spans
@pytest.mark.parametrize("backend", ["asyncio", "trio"])
@pytest.mark.parametrize("sync", [False, True])
def test_background_task_spans(telemetry, backend, sync):
config, exporter, reader = telemetry
app = FastAPI(telemetry=config)
executed = []
endpoint_thread = []
def execute(number, *, value):
span = trace.get_current_span()
assert isinstance(span, Span)
assert span.is_recording()
assert span.name == "fastapi.background_task"
data = get_telemetry_data()
assert data is not None
assert data.request is not None
assert data.request.url.path == "/items/42"
assert len(server_spans(exporter)) == 1
executed.append((number, value, threading.get_ident(), span.get_span_context()))
async def async_task(number, *, value):
execute(number, value=value)
task = execute if sync else async_task
async def dependency(background: BackgroundTasks):
background.add_task(task, 1, value="private dependency value")
@app.get("/items/{item_id}", dependencies=[Depends(dependency)])
async def endpoint(background: BackgroundTasks, item_id: int):
endpoint_thread.append(threading.get_ident())
background.add_task(task, 2, value="private endpoint value")
return {"item_id": item_id}
assert TestClient(app, backend=backend).get("/items/42").json() == {"item_id": 42}
assert [(number, value) for number, value, _, _ in executed] == [
(1, "private dependency value"),
(2, "private endpoint value"),
]
spans = exporter.get_finished_spans()
(server,) = server_spans(exporter)
tasks = [span for span in spans if span.name == "fastapi.background_task"]
assert len(tasks) == 2
for span, (_, _, thread_id, context) in zip(tasks, executed, strict=True):
assert span.kind == SpanKind.INTERNAL
assert span.parent.span_id == server.context.span_id
assert context.trace_id == server.context.trace_id
assert context.span_id == span.context.span_id
assert span.start_time >= server.end_time
assert span.status.status_code == StatusCode.UNSET
assert span.attributes == {
"code.function.name": f"{task.__module__}.{task.__qualname__}"
}
assert (thread_id != endpoint_thread[0]) == sync
assert metric_points(reader=reader)[0].count == 1
assert get_telemetry_data() is None
def test_background_task_objects(telemetry):
config, exporter, _ = telemetry
app = FastAPI(telemetry=config)
executed = []
def record(value):
span = trace.get_current_span()
assert isinstance(span, Span)
executed.append((value, span.name))
class AsyncTask:
async def __call__(self, *, value):
record(value)
class CustomTask(BackgroundTask):
async def __call__(self):
record("before")
await super().__call__()
record("after")
original = CustomTask(record, "custom")
background = BackgroundTasks([original])
background.add_task(partial(record, "partial"))
background.add_task(AsyncTask(), value="async callable")
@app.get("/")
async def endpoint():
return JSONResponse("ok", background=background)
assert TestClient(app).get("/").json() == "ok"
assert background.tasks[0] is original
assert executed == [
(value, "fastapi.background_task")
for value in ("before", "custom", "after", "partial", "async callable")
]
assert (
len(
[
s
for s in exporter.get_finished_spans()
if s.name == "fastapi.background_task"
]
)
== 3
)
@pytest.mark.parametrize("sync", [False, True])
@pytest.mark.parametrize("http_error", [False, True])
def test_background_task_failure(telemetry, logs, sync, http_error):
config, exporter, _ = telemetry
app = FastAPI(telemetry=config)
error = HTTPException(503) if http_error else ValueError("background failure")
executed = []
def fail():
executed.append("first")
raise error
async def async_fail():
fail()
@app.get("/")
async def endpoint(background: BackgroundTasks):
background.add_task(fail if sync else async_fail)
background.add_task(executed.append, "second")
return "ok"
assert TestClient(app, raise_server_exceptions=False).get("/").status_code == 200
assert executed == ["first"]
(task,) = [
s for s in exporter.get_finished_spans() if s.name == "fastapi.background_task"
]
assert task.status.status_code == StatusCode.ERROR
assert task.attributes["error.type"] == (
"fastapi.exceptions.HTTPException" if http_error else "ValueError"
)
assert not task.events
(server,) = server_spans(exporter)
assert server.attributes["http.response.status_code"] == 200
assert server.status.status_code == StatusCode.UNSET
(record,) = logs.get_finished_logs()
assert (
record.log_record.trace_id == task.context.trace_id == server.context.trace_id
)
assert "raise error" in record.log_record.attributes["exception.stacktrace"]
assert record.log_record.exception is not None
if http_error:
assert record.log_record.exception.__cause__ is error
else:
assert record.log_record.exception is error
@pytest.mark.parametrize(
"mode", ["operations_disabled", "tracing_disabled", "excluded", "unsampled"]
)
def test_background_task_settings(telemetry, mode):
config, exporter, _ = telemetry
if mode == "operations_disabled":
config["operation_spans"] = False
elif mode == "tracing_disabled":
config["tracing"] = False
elif mode == "excluded":
config["exclude"] = lambda scope: True
else:
config["tracer_provider"].sampler = ALWAYS_OFF
app = FastAPI(telemetry=config)
executed = []
@app.get("/")
async def endpoint(background: BackgroundTasks):
background.add_task(executed.append, "done")
return "ok"
assert TestClient(app).get("/").json() == "ok"
assert executed == ["done"]
assert not any(
s.name == "fastapi.background_task" for s in exporter.get_finished_spans()
)
def test_background_tasks_outside_request():
executed = []
background = BackgroundTasks()
background.add_task(executed.append, "done")
anyio.run(background)
assert executed == ["done"]
assert get_telemetry_data() is None
@pytest.mark.parametrize("integration", ["contrib", "logfire", "native_logfire"])
@run_in_subprocess
def test_background_task_integrations(integration):
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
InMemorySpanExporter,
)
native = InMemorySpanExporter()
legacy = InMemorySpanExporter()
provider = TracerProvider()
provider.add_span_processor(SimpleSpanProcessor(native))
app = FastAPI(
telemetry={
"tracer_provider": None if integration == "native_logfire" else provider
}
)
executed = []
def sync_task():
executed.append(trace.get_current_span().get_span_context())
async def async_task():
executed.append(trace.get_current_span().get_span_context())
@app.get("/")
async def endpoint(background: BackgroundTasks):
background.add_task(sync_task)
background.add_task(async_task)
return "ok"
if integration == "contrib":
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
legacy_provider = TracerProvider()
legacy_provider.add_span_processor(SimpleSpanProcessor(legacy))
FastAPIInstrumentor.instrument_app(app, tracer_provider=legacy_provider)
else:
import logfire
logfire.configure(
send_to_logfire=False,
console=False,
metrics=False,
additional_span_processors=[SimpleSpanProcessor(legacy)],
)
if integration == "logfire":
logfire.instrument_fastapi(app)
assert TestClient(app).get("/").json() == "ok"
assert not native.get_finished_spans()
tasks = [
span
for span in legacy.get_finished_spans()
if (
span.name.startswith("BackgroundTask ")
or span.name == "fastapi.background_task"
)
and not (
span.attributes
and span.attributes.get("logfire.span_type") == "pending_span"
)
]
assert [span.name for span in tasks] == (
["fastapi.background_task", "fastapi.background_task"]
if integration == "native_logfire"
else ["BackgroundTask sync_task", "BackgroundTask async_task"]
)
assert [span.context for span in tasks] == executed