From bbbb89de42a6f1f3cc813230b80f7e63f9f6e8cd Mon Sep 17 00:00:00 2001
From: Anton Pirker <anton.pirker@sentry.io>
Date: Wed, 4 Dec 2024 09:03:48 +0100
Subject: [PATCH] Fixed tests for boto3

---
 sentry_sdk/integrations/boto3.py    | 13 ++++++++-----
 tests/integrations/boto3/test_s3.py | 12 ++++++------
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/sentry_sdk/integrations/boto3.py b/sentry_sdk/integrations/boto3.py
index dfea7459c3..04eac8c84f 100644
--- a/sentry_sdk/integrations/boto3.py
+++ b/sentry_sdk/integrations/boto3.py
@@ -116,20 +116,19 @@ def _sentry_after_call(context, parsed, **kwargs):
         data=span_data,
     )
 
-    span.__exit__(None, None, None)
-
     body = parsed.get("Body")
     if not isinstance(body, StreamingBody):
+        span.__exit__(None, None, None)
         return
 
-    streaming_span = span.start_child(
+    streaming_span = sentry_sdk.start_span(
         op=OP.HTTP_CLIENT_STREAM,
-        name=span.description,
+        name=span.name,
         origin=Boto3Integration.origin,
+        only_if_parent=True,
     )
 
     orig_read = body.read
-    orig_close = body.close
 
     def sentry_streaming_body_read(*args, **kwargs):
         # type: (*Any, **Any) -> bytes
@@ -144,6 +143,8 @@ def sentry_streaming_body_read(*args, **kwargs):
 
     body.read = sentry_streaming_body_read
 
+    orig_close = body.close
+
     def sentry_streaming_body_close(*args, **kwargs):
         # type: (*Any, **Any) -> None
         streaming_span.finish()
@@ -151,6 +152,8 @@ def sentry_streaming_body_close(*args, **kwargs):
 
     body.close = sentry_streaming_body_close
 
+    span.__exit__(None, None, None)
+
 
 def _sentry_after_call_error(context, exception, **kwargs):
     # type: (Dict[str, Any], Type[BaseException], **Any) -> None
diff --git a/tests/integrations/boto3/test_s3.py b/tests/integrations/boto3/test_s3.py
index 668e8349b6..71dc5ccc07 100644
--- a/tests/integrations/boto3/test_s3.py
+++ b/tests/integrations/boto3/test_s3.py
@@ -21,7 +21,7 @@ def test_basic(sentry_init, capture_events):
     events = capture_events()
 
     s3 = session.resource("s3")
-    with sentry_sdk.start_transaction() as transaction, MockResponse(
+    with sentry_sdk.start_span() as transaction, MockResponse(
         s3.meta.client, 200, {}, read_fixture("s3_list.xml")
     ):
         bucket = s3.Bucket("bucket")
@@ -45,7 +45,7 @@ def test_breadcrumb(sentry_init, capture_events):
 
     try:
         s3 = session.resource("s3")
-        with sentry_sdk.start_transaction(), MockResponse(
+        with sentry_sdk.start_span(), MockResponse(
             s3.meta.client, 200, {}, read_fixture("s3_list.xml")
         ):
             bucket = s3.Bucket("bucket")
@@ -75,7 +75,7 @@ def test_streaming(sentry_init, capture_events):
     events = capture_events()
 
     s3 = session.resource("s3")
-    with sentry_sdk.start_transaction() as transaction, MockResponse(
+    with sentry_sdk.start_span() as transaction, MockResponse(
         s3.meta.client, 200, {}, b"hello"
     ):
         obj = s3.Bucket("bucket").Object("foo.pdf")
@@ -113,7 +113,7 @@ def test_streaming_close(sentry_init, capture_events):
     events = capture_events()
 
     s3 = session.resource("s3")
-    with sentry_sdk.start_transaction() as transaction, MockResponse(
+    with sentry_sdk.start_span() as transaction, MockResponse(
         s3.meta.client, 200, {}, b"hello"
     ):
         obj = s3.Bucket("bucket").Object("foo.pdf")
@@ -142,7 +142,7 @@ def test_omit_url_data_if_parsing_fails(sentry_init, capture_events):
         "sentry_sdk.integrations.boto3.parse_url",
         side_effect=ValueError,
     ):
-        with sentry_sdk.start_transaction() as transaction, MockResponse(
+        with sentry_sdk.start_span() as transaction, MockResponse(
             s3.meta.client, 200, {}, read_fixture("s3_list.xml")
         ):
             bucket = s3.Bucket("bucket")
@@ -170,7 +170,7 @@ def test_span_origin(sentry_init, capture_events):
     events = capture_events()
 
     s3 = session.resource("s3")
-    with sentry_sdk.start_transaction(), MockResponse(
+    with sentry_sdk.start_span(), MockResponse(
         s3.meta.client, 200, {}, read_fixture("s3_list.xml")
     ):
         bucket = s3.Bucket("bucket")