diff --git a/elasticapm/contrib/opentelemetry/span.py b/elasticapm/contrib/opentelemetry/span.py index a7d049faa..7917070d5 100644 --- a/elasticapm/contrib/opentelemetry/span.py +++ b/elasticapm/contrib/opentelemetry/span.py @@ -32,7 +32,7 @@ import types as python_types import typing import urllib.parse -from typing import Optional +from typing import Optional, Union from opentelemetry.context import Context from opentelemetry.sdk import trace as oteltrace @@ -157,13 +157,19 @@ def is_recording(self) -> bool: """ return self.elastic_span.transaction.is_sampled and not self.elastic_span.ended_time - def set_status(self, status: Status) -> None: + def set_status(self, status: Union[Status, StatusCode], description: Optional[str] = None) -> None: """Sets the Status of the Span. If used, this will override the default Span status. """ - if status.status_code == StatusCode.ERROR: + # Handle both Status objects and StatusCode enums + if isinstance(status, Status): + status_code = status.status_code + else: + status_code = status + + if status_code == StatusCode.ERROR: self.elastic_span.outcome = constants.OUTCOME.FAILURE - elif status.status_code == StatusCode.OK: + elif status_code == StatusCode.OK: self.elastic_span.outcome = constants.OUTCOME.SUCCESS else: self.elastic_span.outcome = constants.OUTCOME.UNKNOWN diff --git a/tests/contrib/opentelemetry/tests.py b/tests/contrib/opentelemetry/tests.py index 3a302289e..aa0160b74 100755 --- a/tests/contrib/opentelemetry/tests.py +++ b/tests/contrib/opentelemetry/tests.py @@ -34,6 +34,7 @@ from opentelemetry.trace import Link, SpanContext, SpanKind, TraceFlags from opentelemetry.trace.propagation import _SPAN_KEY +from opentelemetry.trace.status import Status, StatusCode import elasticapm.contrib.opentelemetry.context as context import elasticapm.contrib.opentelemetry.trace as trace @@ -155,4 +156,51 @@ def test_span_links(tracer: Tracer): assert span["links"][0]["span_id"] == "0011223344556677" +def test_set_status_with_status_object(tracer: Tracer): + with tracer.start_as_current_span("test") as span: + span.set_status(Status(StatusCode.OK)) + + client = tracer.client + transaction = client.events[constants.TRANSACTION][0] + assert transaction["outcome"] == "success" + + +def test_set_status_with_status_code(tracer: Tracer): + with tracer.start_as_current_span("test") as span: + span.set_status(StatusCode.ERROR) + + client = tracer.client + transaction = client.events[constants.TRANSACTION][0] + assert transaction["outcome"] == "failure" + + +def test_set_status_with_status_code_and_description(tracer: Tracer): + with tracer.start_as_current_span("test") as span: + span.set_status(StatusCode.OK, "Everything is fine") + + client = tracer.client + transaction = client.events[constants.TRANSACTION][0] + assert transaction["outcome"] == "success" + + +def test_set_status_unset(tracer: Tracer): + with tracer.start_as_current_span("test") as span: + span.set_status(StatusCode.UNSET) + + client = tracer.client + transaction = client.events[constants.TRANSACTION][0] + assert transaction["outcome"] == "unknown" + + +def test_set_status_on_span(tracer: Tracer): + """Test set_status on a child span (not transaction)""" + with tracer.start_as_current_span("test"): + with tracer.start_as_current_span("testspan") as span: + span.set_status(StatusCode.ERROR) + + client = tracer.client + span_event = client.events[constants.SPAN][0] + assert span_event["outcome"] == "failure" + + # TODO Add some span subtype testing?