refactor: extract set_send_span and set_server_span logic to helpers

Author: omgitsaheadcrab

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

@@ -829,6 +829,66 @@ async def otel_receive():
 
         return otel_receive
 
+    def _set_send_span(
+        self,
+        server_span_name,
+        scope,
+        send,
+        message,
+        status_code,
+        expecting_trailers,
+    ):
+        """Set send span attributes and status code."""
+        with self.tracer.start_as_current_span(
+            " ".join((server_span_name, scope["type"], "send"))
+        ) as send_span:
+            if callable(self.client_response_hook):
+                self.client_response_hook(send_span, scope, message)
+
+            if send_span.is_recording():
+                if message["type"] == "http.response.start":
+                    expecting_trailers = message.get("trailers", False)
+                send_span.set_attribute("asgi.event.type", message["type"])
+
+            if status_code:
+                set_status_code(
+                    send_span,
+                    status_code,
+                    None,
+                    self._sem_conv_opt_in_mode,
+                )
+        return expecting_trailers
+
+    def _set_server_span(
+        self, server_span, message, status_code, duration_attrs
+    ):
+        """Set server span attributes and status code."""
+        if (
+            server_span.is_recording()
+            and server_span.kind == trace.SpanKind.SERVER
+            and "headers" in message
+        ):
+            custom_response_attributes = (
+                collect_custom_headers_attributes(
+                    message,
+                    self.http_capture_headers_sanitize_fields,
+                    self.http_capture_headers_server_response,
+                    normalise_response_header_name,
+                )
+                if self.http_capture_headers_server_response
+                else {}
+            )
+            if len(custom_response_attributes) > 0:
+                server_span.set_attributes(custom_response_attributes)
+
+        if status_code:
+            set_status_code(
+                server_span,
+                status_code,
+                duration_attrs,
+                self._sem_conv_opt_in_mode,
+            )
+
     def _get_otel_send(
         self,
         server_span,
@@ -850,52 +910,19 @@ async def otel_send(message: dict[str, Any]):
                 status_code = 200
 
             if not self.exclude_send_span:
-                with self.tracer.start_as_current_span(
-                    " ".join((server_span_name, scope["type"], "send"))
-                ) as send_span:
-                    if callable(self.client_response_hook):
-                        self.client_response_hook(send_span, scope, message)
-
-                    if send_span.is_recording():
-                        if message["type"] == "http.response.start":
-                            expecting_trailers = message.get("trailers", False)
-                        send_span.set_attribute(
-                            "asgi.event.type", message["type"]
-                        )
-                        if status_code:
-                            set_status_code(
-                                send_span,
-                                status_code,
-                                None,
-                                self._sem_conv_opt_in_mode,
-                            )
-
-            if (
-                server_span.is_recording()
-                and server_span.kind == trace.SpanKind.SERVER
-                and "headers" in message
-            ):
-                custom_response_attributes = (
-                    collect_custom_headers_attributes(
-                        message,
-                        self.http_capture_headers_sanitize_fields,
-                        self.http_capture_headers_server_response,
-                        normalise_response_header_name,
-                    )
-                    if self.http_capture_headers_server_response
-                    else {}
-                )
-                if len(custom_response_attributes) > 0:
-                    server_span.set_attributes(custom_response_attributes)
-
-            if status_code:
-                set_status_code(
-                    server_span,
+                expecting_trailers = self._set_send_span(
+                    server_span_name,
+                    scope,
+                    send,
+                    message,
                     status_code,
-                    duration_attrs,
-                    self._sem_conv_opt_in_mode,
+                    expecting_trailers,
                 )
 
+            self._set_server_span(
+                server_span, message, status_code, duration_attrs
+            )
+
             propagator = get_global_response_propagator()
             if propagator:
                 propagator.inject(