Return inner span processor method results

DiamondLightSource · Apr 18, 2024 · 2a4caa9 · 2a4caa9
1 parent d186a95
commit 2a4caa9
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/src/graph_energy_scan/telemetry.py b/src/graph_energy_scan/telemetry.py
@@ -18,17 +18,20 @@ def __init__(self, inner: SpanProcessor, filter_types: list[str]) -> None:
         super().__init__()
 
     def on_start(self, span: Span, parent_context: trace.Context | None = None) -> None:
-        self.inner.on_start(span, parent_context)
+        return self.inner.on_start(span, parent_context)
 
     def on_end(self, span: ReadableSpan) -> None:
-        if span.attributes.get("type") not in self.filter_types:
-            self.inner.on_end(span)
+        if (
+            span.attributes is not None
+            and span.attributes.get("type") not in self.filter_types
+        ):
+            return self.inner.on_end(span)
 
     def shutdown(self) -> None:
-        self.inner.shutdown()
+        return self.inner.shutdown()
 
     def force_flush(self, timeout_millis: int = 30000) -> bool:
-        self.inner.force_flush(timeout_millis)
+        return self.inner.force_flush(timeout_millis)
 
 
 def setup_telemetry(otel_collector_url: str):