update to Observability-utils 0.1.4 (#705)

dev-requirements.txt

@@ -123,7 +123,7 @@ mongoquery==1.4.2
 msgpack==1.1.0
 msgpack-numpy==0.4.8
 multidict==6.1.0
-mypy==1.13.0
+mypy==1.13.0 
 mypy-extensions==1.0.0
 myst-parser==4.0.0
 networkx==3.4.2
@@ -132,7 +132,7 @@ nose2==0.15.1
 nslsii==0.10.5
 numcodecs==0.13.1
 numpy==1.26.4
-observability-utils==0.1.2
+observability-utils==0.1.4
 opencv-python-headless==4.10.0.84
 opentelemetry-api==1.27.0
 opentelemetry-distro==0.48b0

pyproject.toml

@@ -35,7 +35,7 @@ dependencies = [
     "event-model==1.21",    # https://github.com/DiamondLightSource/blueapi/issues/684
     "opentelemetry-distro>=0.48b0",
     "opentelemetry-instrumentation-fastapi>=0.48b0",
-    "observability-utils>=0.1.2",
+    "observability-utils>=0.1.4"
 ]
 dynamic = ["version"]
 license.file = "LICENSE"

src/blueapi/worker/task_worker.py

@@ -11,13 +11,12 @@
 from bluesky.protocols import Status
 from observability_utils.tracing import (
     add_span_attributes,
-    get_trace_context,
     get_tracer,
     setup_tracing,
     start_as_current_span,
 )
 from opentelemetry.baggage import get_baggage
-from opentelemetry.context import Context
+from opentelemetry.context import Context, get_current
 from opentelemetry.trace import SpanKind
 from pydantic import Field
 from super_state_machine.errors import TransitionError
@@ -222,7 +221,7 @@ def mark_task_as_started(event: WorkerEvent, _: str | None) -> None:
 
         LOGGER.info(f"Submitting: {trackable_task}")
         try:
-            self._current_task_otel_context = get_trace_context()
+            self._current_task_otel_context = get_current()
             sub = self.worker_events.subscribe(mark_task_as_started)
             """ Cache the current trace context as the one for this task id """
             self._task_channel.put_nowait(trackable_task)
@@ -326,7 +325,7 @@ def _cycle(self) -> None:
                             next_task  # Informing mypy that the task is not None
                         )
 
-                        self._current_task_otel_context = get_trace_context()
+                        self._current_task_otel_context = get_current()
                         add_span_attributes({"next_task.task_id": next_task.task_id})
 
                         self._current.is_pending = False

tests/conftest.py

@@ -5,13 +5,11 @@
 import pytest
 from bluesky import RunEngine
 from bluesky.run_engine import TransitionError
-from observability_utils.tracing import setup_tracing
+from observability_utils.tracing import JsonObjectSpanExporter, setup_tracing
 from opentelemetry.sdk.trace import TracerProvider
-from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.trace import get_tracer_provider
 
-from tests.unit_tests.utils.test_tracing import JsonObjectSpanExporter
-
 
 @pytest.fixture(scope="function")
 def RE(request):
@@ -33,15 +31,11 @@ def clean_event_loop():
     return RE
 
 
-@pytest.fixture
-def provider() -> TracerProvider:
+@pytest.fixture(scope="session")
+def exporter() -> TracerProvider:
     setup_tracing("test", False)
-    return cast(TracerProvider, get_tracer_provider())
-
-
-@pytest.fixture
-def exporter(provider: TracerProvider) -> JsonObjectSpanExporter:
     exporter = JsonObjectSpanExporter()
-    processor = BatchSpanProcessor(exporter)
-    provider.add_span_processor(processor)
+    provider = cast(TracerProvider, get_tracer_provider())
+    # Use SimpleSpanProcessor to keep tests quick
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
     return exporter

tests/unit_tests/client/test_client.py

@@ -3,7 +3,7 @@
 
 import pytest
 from bluesky_stomp.messaging import MessageContext
-from tests.unit_tests.utils.test_tracing import (
+from observability_utils.tracing import (
     JsonObjectSpanExporter,
     asserting_span_exporter,
 )

tests/unit_tests/service/test_runner.py

@@ -3,12 +3,12 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
-from ophyd import Callable
-from pydantic import BaseModel, ValidationError
-from tests.unit_tests.utils.test_tracing import (
+from observability_utils.tracing import (
     JsonObjectSpanExporter,
     asserting_span_exporter,
 )
+from ophyd import Callable
+from pydantic import BaseModel, ValidationError
 
 from blueapi.service import interface
 from blueapi.service.model import EnvironmentResponse

tests/unit_tests/worker/test_task_worker.py

@@ -7,7 +7,7 @@
 from unittest.mock import ANY, MagicMock, patch
 
 import pytest
-from tests.unit_tests.utils.test_tracing import (
+from observability_utils.tracing import (
     JsonObjectSpanExporter,
     asserting_span_exporter,
 )
@@ -554,3 +554,11 @@ def test_clear_task_span_ok(
     with pytest.raises(KeyError):
         with asserting_span_exporter(exporter, "clear_task", "task_id"):
             worker.clear_task("foo")
+
+
+def test_begin_task_span_ok(
+    exporter: JsonObjectSpanExporter, worker: TaskWorker
+) -> None:
+    task_id = worker.submit_task(_SIMPLE_TASK)
+    with asserting_span_exporter(exporter, "begin_task", "task_id"):
+        worker.begin_task(task_id)