Files

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