diff --git a/taskiq/instrumentation.py b/taskiq/instrumentation.py index ed542c7b..51b9347b 100644 --- a/taskiq/instrumentation.py +++ b/taskiq/instrumentation.py @@ -87,7 +87,7 @@ def instrument_broker( if not hasattr(broker, "_is_instrumented_by_opentelemetry"): broker._is_instrumented_by_opentelemetry = False # type: ignore[attr-defined] # noqa: SLF001 - if not getattr(broker, "is_instrumented_by_opentelemetry", False): + if not getattr(broker, "_is_instrumented_by_opentelemetry", False): broker.middlewares.insert( 0, OpenTelemetryMiddleware( diff --git a/tests/opentelemetry/test_auto_instrumentation.py b/tests/opentelemetry/test_auto_instrumentation.py index c7290869..305d5068 100644 --- a/tests/opentelemetry/test_auto_instrumentation.py +++ b/tests/opentelemetry/test_auto_instrumentation.py @@ -8,6 +8,10 @@ class TestTaskiqAutoInstrumentation(TestBase): + def tearDown(self) -> None: + TaskiqInstrumentor().uninstrument() + super().tearDown() + def test_auto_instrument(self) -> None: TaskiqInstrumentor().instrument() diff --git a/tests/opentelemetry/test_duplicate.py b/tests/opentelemetry/test_duplicate.py index b03ff8fc..e0f84e03 100644 --- a/tests/opentelemetry/test_duplicate.py +++ b/tests/opentelemetry/test_duplicate.py @@ -9,7 +9,10 @@ class TestUtils(unittest.TestCase): @staticmethod def test_duplicate_instrumentation() -> None: first = TaskiqInstrumentor() - first.instrument_broker(broker) - second = TaskiqInstrumentor() - second.instrument_broker(broker) - TaskiqInstrumentor().uninstrument() + try: + first.instrument_broker(broker) + second = TaskiqInstrumentor() + second.instrument_broker(broker) + assert len(broker.middlewares) == 1 + finally: + first.uninstrument_broker(broker) diff --git a/tests/opentelemetry/test_idempotence.py b/tests/opentelemetry/test_idempotence.py new file mode 100644 index 00000000..f3ad6a23 --- /dev/null +++ b/tests/opentelemetry/test_idempotence.py @@ -0,0 +1,87 @@ +from collections.abc import Generator + +import pytest +from opentelemetry import context +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import SpanKind + +from taskiq import InMemoryBroker +from taskiq.instrumentation import TaskiqInstrumentor +from taskiq.middlewares.opentelemetry_middleware import OpenTelemetryMiddleware + + +@pytest.fixture +def local_broker() -> Generator[InMemoryBroker, None, None]: + broker = InMemoryBroker(await_inplace=True) + TaskiqInstrumentor().uninstrument_broker(broker) + try: + yield broker + finally: + TaskiqInstrumentor().uninstrument_broker(broker) + broker.executor.shutdown() + + +def test_repeated_instrumentation_keeps_one_middleware( + local_broker: InMemoryBroker, +) -> None: + instrumentor = TaskiqInstrumentor() + instrumentor.instrument_broker(local_broker) + first = local_broker.middlewares[0] + instrumentor.instrument_broker(local_broker) + assert local_broker.middlewares == [first] + + +def test_uninstrumented_broker_can_be_instrumented_again( + local_broker: InMemoryBroker, +) -> None: + instrumentor = TaskiqInstrumentor() + instrumentor.instrument_broker(local_broker) + instrumentor.uninstrument_broker(local_broker) + assert local_broker.middlewares == [] + instrumentor.instrument_broker(local_broker) + assert len(local_broker.middlewares) == 1 + assert isinstance(local_broker.middlewares[0], OpenTelemetryMiddleware) + + +def test_different_brokers_are_instrumented_independently( + local_broker: InMemoryBroker, +) -> None: + other = InMemoryBroker() + instrumentor = TaskiqInstrumentor() + try: + instrumentor.instrument_broker(local_broker) + instrumentor.instrument_broker(other) + assert len(local_broker.middlewares) == len(other.middlewares) == 1 + assert local_broker.middlewares[0] is not other.middlewares[0] + finally: + instrumentor.uninstrument_broker(other) + other.executor.shutdown() + + +async def test_repeated_setup_does_not_duplicate_task_spans( + local_broker: InMemoryBroker, +) -> None: + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + instrumentor = TaskiqInstrumentor() + token = context.attach(context.Context()) + try: + instrumentor.instrument_broker(local_broker, tracer_provider=provider) + instrumentor.instrument_broker(local_broker, tracer_provider=provider) + + @local_broker.task + async def task() -> int: + return 42 + + result = await task.kiq() + assert (await result.get_result()).return_value == 42 + spans = exporter.get_finished_spans() + assert len([span for span in spans if span.kind == SpanKind.PRODUCER]) == 1 + assert len([span for span in spans if span.kind == SpanKind.CONSUMER]) == 1 + finally: + instrumentor.uninstrument_broker(local_broker) + context.detach(token) + provider.shutdown()