fix(contrib/opentelemetry): set_status should match base signature (#2457)

* fix(contrib/opentelemetry): set_status should match base signature

fixes #2455

Signed-off-by: qvalentin <[email protected]>

* Apply suggestions from code review

---------

Signed-off-by: qvalentin <[email protected]>
Co-authored-by: Riccardo Magliocchetti <[email protected]>

Author: qvalentin

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

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?