Set status in start_as_current_span too (#381)

Fixes #377
open-telemetry · Jan 30, 2020 · 31f5726 · 31f5726
1 parent 6aa69ae
commit 31f5726
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 10 deletions.
diff --git a/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py b/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py
@@ -538,6 +538,23 @@ def use_span(
                 self.source._current_span_slot.set(  # pylint:disable=protected-access
                     span_snapshot
                 )
+
+        except Exception as error:  # pylint: disable=broad-except
+            if (
+                span.status is None
+                and span._set_status_on_exception  # pylint:disable=protected-access  # noqa
+            ):
+                span.set_status(
+                    Status(
+                        canonical_code=StatusCanonicalCode.UNKNOWN,
+                        description="{}: {}".format(
+                            type(error).__name__, error
+                        ),
+                    )
+                )
+
+                raise
+
         finally:
             if end_on_exit:
                 span.end()

diff --git a/opentelemetry-sdk/tests/trace/test_trace.py b/opentelemetry-sdk/tests/trace/test_trace.py
@@ -641,16 +641,26 @@ def test_ended_span(self):
         )
 
     def test_error_status(self):
-        try:
-            with trace.TracerSource().get_tracer(__name__).start_span(
-                "root"
-            ) as root:
-                raise Exception("unknown")
-        except Exception:  # pylint: disable=broad-except
-            pass
-
-        self.assertIs(root.status.canonical_code, StatusCanonicalCode.UNKNOWN)
-        self.assertEqual(root.status.description, "Exception: unknown")
+        def error_status_test(context):
+            with self.assertRaises(AssertionError):
+                with context as root:
+                    raise AssertionError("unknown")
+
+            self.assertIs(
+                root.status.canonical_code, StatusCanonicalCode.UNKNOWN
+            )
+            self.assertEqual(
+                root.status.description, "AssertionError: unknown"
+            )
+
+        error_status_test(
+            trace.TracerSource().get_tracer(__name__).start_span("root")
+        )
+        error_status_test(
+            trace.TracerSource()
+            .get_tracer(__name__)
+            .start_as_current_span("root")
+        )
 
 
 def span_event_start_fmt(span_processor_name, span_name):