Refactor metrics_context assignment, update test

Author: tammy-baylis-swi

instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py

@@ -290,7 +290,6 @@ def response_hook(span: Span, status: str, response_headers: List):
 from opentelemetry.semconv.metrics.http_metrics import (
     HTTP_SERVER_REQUEST_DURATION,
 )
-from opentelemetry.trace.propagation import _SPAN_KEY
 from opentelemetry.util._importlib_metadata import version
 from opentelemetry.util.http import (
     get_excluded_urls,
@@ -412,6 +411,11 @@ def _start_response(status, response_headers, *args, **kwargs):
         result = wsgi_app(wrapped_app_environ, _start_response)
         if should_trace:
             duration_s = default_timer() - start
+            # Get the span from wrapped_app_environ and re-create context manually
+            # to pass to histogram for exemplars generation
+            span = wrapped_app_environ.get(_ENVIRON_SPAN_KEY)
+            metrics_context = trace.set_span_in_context(span)
+
             if duration_histogram_old:
                 duration_attrs_old = otel_wsgi._parse_duration_attrs(
                     attributes, _StabilityMode.DEFAULT
@@ -420,13 +424,6 @@ def _start_response(status, response_headers, *args, **kwargs):
                 if request_route:
                     # http.target to be included in old semantic conventions
                     duration_attrs_old[HTTP_TARGET] = str(request_route)
-
-                # Get the span from wrapped_app_environ and re-create context manually
-                # to pass to histogram for exemplars generation
-                span = wrapped_app_environ.get(_ENVIRON_SPAN_KEY)
-                metrics_context = (
-                    context.set_value(_SPAN_KEY, span) if span else None
-                )
                 duration_histogram_old.record(
                     max(round(duration_s * 1000), 0),
                     duration_attrs_old,
@@ -440,12 +437,6 @@ def _start_response(status, response_headers, *args, **kwargs):
                 if request_route:
                     duration_attrs_new[HTTP_ROUTE] = str(request_route)
 
-                # Get the span from wrapped_app_environ and re-create context manually
-                # to pass to histogram for exemplars generation
-                span = wrapped_app_environ.get(_ENVIRON_SPAN_KEY)
-                metrics_context = (
-                    context.set_value(_SPAN_KEY, span) if span else None
-                )
                 duration_histogram_new.record(
                     max(duration_s, 0),
                     duration_attrs_new,

instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py

@@ -69,7 +69,6 @@
 )
 from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
 from opentelemetry.test.wsgitestutil import WsgiTestBase
-from opentelemetry.trace.propagation import _SPAN_KEY
 from opentelemetry.util.http import (
     OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS,
     OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST,
@@ -812,16 +811,15 @@ def test_flask_metrics_excluded_urls_new_semconv(self):
         self.assertFalse(histogram_data_point_seen)
 
     def test_duration_histogram_old_record_with_context(self):
-        with patch("opentelemetry.context.set_value") as mock_set_value:
+        with patch("opentelemetry.trace.set_span_in_context") as mock_set_span:
             self.client.get("/hello/123")
 
-            # Verify that context.set_value was called for metrics exemplar context
+            # Verify that trace.set_span_in_context was called for metrics exemplar context
             # with same trace, span ID as trace
-            mock_set_value.assert_called()
-            call_args = mock_set_value.call_args
-            self.assertEqual(len(call_args[0]), 2)
-            self.assertEqual(call_args[0][0], _SPAN_KEY)
-            span_arg = call_args[0][1]
+            mock_set_span.assert_called()
+            call_args = mock_set_span.call_args
+            self.assertEqual(len(call_args[0]), 1)
+            span_arg = call_args[0][0]
             self.assertIsNotNone(span_arg)
             finished_spans = self.memory_exporter.get_finished_spans()
             self.assertEqual(len(finished_spans), 1)
@@ -834,16 +832,15 @@ def test_duration_histogram_old_record_with_context(self):
             )
 
     def test_duration_histogram_new_record_with_context_new_semconv(self):
-        with patch("opentelemetry.context.set_value") as mock_set_value:
+        with patch("opentelemetry.trace.set_span_in_context") as mock_set_span:
             self.client.get("/hello/123")
 
-            # Verify that context.set_value was called for metrics exemplar context
+            # Verify that trace.set_span_in_context was called for metrics exemplar context
             # with same trace, span ID as trace
-            mock_set_value.assert_called()
-            call_args = mock_set_value.call_args
-            self.assertEqual(len(call_args[0]), 2)
-            self.assertEqual(call_args[0][0], _SPAN_KEY)
-            span_arg = call_args[0][1]
+            mock_set_span.assert_called()
+            call_args = mock_set_span.call_args
+            self.assertEqual(len(call_args[0]), 1)
+            span_arg = call_args[0][0]
             self.assertIsNotNone(span_arg)
             finished_spans = self.memory_exporter.get_finished_spans()
             self.assertEqual(len(finished_spans), 1)