test(flask): Add span streaming test coverage (#6264)

Add test coverage to ensure that span streaming works for Flask. No
additional changes are needed because we are no longer adding the
request body to the streamed spans, and span streaming support has been
added earlier to the uWSGI middleware which the Flask integration
leverages.

Fixes PY-2323
Fixes #6021

Author: ericapisani

sentry_sdk/integrations/flask.py

@@ -95,10 +95,9 @@ def setup_once() -> None:
         def sentry_patched_wsgi_app(
             self: "Any", environ: "Dict[str, str]", start_response: "Callable[..., Any]"
         ) -> "_ScopedResponse":
-            if sentry_sdk.get_client().get_integration(FlaskIntegration) is None:
-                return old_app(self, environ, start_response)
-
             integration = sentry_sdk.get_client().get_integration(FlaskIntegration)
+            if integration is None:
+                return old_app(self, environ, start_response)
 
             middleware = SentryWsgiMiddleware(
                 lambda *a, **kw: old_app(self, *a, **kw),

tests/integrations/flask/test_flask.py

@@ -15,6 +15,8 @@
 from flask.views import View
 from flask_login import LoginManager, login_user
 
+from sentry_sdk.traces import SpanStatus
+
 try:
     from werkzeug.wrappers.request import UnsupportedMediaType
 except ImportError:
@@ -82,6 +84,7 @@ def test_has_context(sentry_init, app, capture_events):
     assert event["request"]["url"] == "http://localhost/message"
 
 
+@pytest.mark.parametrize("span_streaming", [True, False])
 @pytest.mark.parametrize(
     "url,transaction_style,expected_transaction,expected_source",
     [
@@ -91,29 +94,45 @@ def test_has_context(sentry_init, app, capture_events):
         ("/message/123456", "url", "/message/<int:message_id>", "route"),
     ],
 )
-def test_transaction_style(
+def test_transaction_or_segment_style(
     sentry_init,
     app,
     capture_events,
+    capture_items,
     url,
     transaction_style,
     expected_transaction,
     expected_source,
+    span_streaming,
 ):
     sentry_init(
         integrations=[
             flask_sentry.FlaskIntegration(transaction_style=transaction_style)
-        ]
+        ],
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream" if span_streaming else "static"},
     )
-    events = capture_events()
+
+    if span_streaming:
+        items = capture_items("span")
+    else:
+        events = capture_events()
 
     client = app.test_client()
     response = client.get(url)
     assert response.status_code == 200
 
-    (event,) = events
-    assert event["transaction"] == expected_transaction
-    assert event["transaction_info"] == {"source": expected_source}
+    if span_streaming:
+        sentry_sdk.flush()
+        spans = [i.payload for i in items if i.type == "span"]
+        assert len(spans) == 1
+        (segment,) = spans
+        assert segment["name"] == expected_transaction
+        assert segment["attributes"]["sentry.span.source"] == expected_source
+    else:
+        (_, event) = events
+        assert event["transaction"] == expected_transaction
+        assert event["transaction_info"] == {"source": expected_source}
 
 
 @pytest.mark.parametrize("debug", (True, False))
@@ -763,8 +782,15 @@ def zerodivision(e):
     assert not events
 
 
-def test_tracing_success(sentry_init, capture_events, app):
-    sentry_init(traces_sample_rate=1.0, integrations=[flask_sentry.FlaskIntegration()])
+@pytest.mark.parametrize("span_streaming", [True, False])
+def test_tracing_success(
+    sentry_init, capture_events, capture_items, app, span_streaming
+):
+    sentry_init(
+        traces_sample_rate=1.0,
+        integrations=[flask_sentry.FlaskIntegration()],
+        _experiments={"trace_lifecycle": "stream" if span_streaming else "static"},
+    )
 
     @app.before_request
     def _():
@@ -776,30 +802,61 @@ def hi_tx():
         capture_message("hi")
         return "ok"
 
-    events = capture_events()
+    if span_streaming:
+        items = capture_items("event", "span")
+    else:
+        events = capture_events()
 
     with app.test_client() as client:
         response = client.get("/message_tx")
         assert response.status_code == 200
 
-    message_event, transaction_event = events
+    if span_streaming:
+        sentry_sdk.flush()
+        spans = [i.payload for i in items if i.type == "span"]
+        message_events = [i.payload for i in items if i.type == "event"]
 
-    assert transaction_event["type"] == "transaction"
-    assert transaction_event["transaction"] == "hi_tx"
-    assert transaction_event["contexts"]["trace"]["status"] == "ok"
-    assert transaction_event["tags"]["view"] == "yes"
-    assert transaction_event["tags"]["before_request"] == "yes"
+        assert len(spans) == 1
+        assert len(message_events) == 1
 
-    assert message_event["message"] == "hi"
-    assert message_event["transaction"] == "hi_tx"
-    assert message_event["tags"]["view"] == "yes"
-    assert message_event["tags"]["before_request"] == "yes"
+        (segment,) = spans
+        (message_event,) = message_events
 
+        assert segment["name"] == "hi_tx"
+        assert segment["status"] == SpanStatus.OK
+        assert segment["attributes"]["sentry.origin"] == "auto.http.flask"
 
-def test_tracing_error(sentry_init, capture_events, app):
-    sentry_init(traces_sample_rate=1.0, integrations=[flask_sentry.FlaskIntegration()])
+        assert message_event["message"] == "hi"
+        assert message_event["transaction"] == "hi_tx"
+        assert message_event["tags"]["view"] == "yes"
+        assert message_event["tags"]["before_request"] == "yes"
+    else:
+        message_event, transaction_event = events
 
-    events = capture_events()
+        assert transaction_event["type"] == "transaction"
+        assert transaction_event["transaction"] == "hi_tx"
+        assert transaction_event["contexts"]["trace"]["status"] == "ok"
+        assert transaction_event["tags"]["view"] == "yes"
+        assert transaction_event["tags"]["before_request"] == "yes"
+
+        assert message_event["message"] == "hi"
+        assert message_event["transaction"] == "hi_tx"
+        assert message_event["tags"]["view"] == "yes"
+        assert message_event["tags"]["before_request"] == "yes"
+
+
+@pytest.mark.parametrize("span_streaming", [True, False])
+def test_tracing_error(sentry_init, capture_events, capture_items, app, span_streaming):
+    sentry_init(
+        traces_sample_rate=1.0,
+        integrations=[flask_sentry.FlaskIntegration()],
+        _experiments={"trace_lifecycle": "stream" if span_streaming else "static"},
+    )
+
+    if span_streaming:
+        items = capture_items("event", "span")
+    else:
+        events = capture_events()
 
     @app.route("/error")
     def error():
@@ -810,15 +867,33 @@ def error():
             response = client.get("/error")
             assert response.status_code == 500
 
-    error_event, transaction_event = events
+    if span_streaming:
+        sentry_sdk.flush()
+        spans = [i.payload for i in items if i.type == "span"]
+        error_events = [i.payload for i in items if i.type == "event"]
 
-    assert transaction_event["type"] == "transaction"
-    assert transaction_event["transaction"] == "error"
-    assert transaction_event["contexts"]["trace"]["status"] == "internal_error"
+        assert len(spans) == 1
+        assert len(error_events) == 1
 
-    assert error_event["transaction"] == "error"
-    (exception,) = error_event["exception"]["values"]
-    assert exception["type"] == "ZeroDivisionError"
+        (segment,) = spans
+        (error_event,) = error_events
+
+        assert segment["name"] == "error"
+        assert segment["status"] == SpanStatus.ERROR
+
+        assert error_event["transaction"] == "error"
+        (exception,) = error_event["exception"]["values"]
+        assert exception["type"] == "ZeroDivisionError"
+    else:
+        error_event, transaction_event = events
+
+        assert transaction_event["type"] == "transaction"
+        assert transaction_event["transaction"] == "error"
+        assert transaction_event["contexts"]["trace"]["status"] == "internal_error"
+
+        assert error_event["transaction"] == "error"
+        (exception,) = error_event["exception"]["values"]
+        assert exception["type"] == "ZeroDivisionError"
 
 
 def test_error_has_trace_context_if_tracing_disabled(sentry_init, capture_events, app):
@@ -995,34 +1070,54 @@ def test_response_status_code_not_found_in_transaction_context(
     assert transaction["contexts"]["response"]["status_code"] == 404
 
 
-def test_span_origin(sentry_init, app, capture_events):
+@pytest.mark.parametrize("span_streaming", [True, False])
+def test_span_origin(sentry_init, app, capture_events, capture_items, span_streaming):
     sentry_init(
         integrations=[flask_sentry.FlaskIntegration()],
         traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream" if span_streaming else "static"},
     )
-    events = capture_events()
+
+    if span_streaming:
+        items = capture_items("span")
+    else:
+        events = capture_events()
 
     client = app.test_client()
     client.get("/message")
 
-    (_, event) = events
-
-    assert event["contexts"]["trace"]["origin"] == "auto.http.flask"
+    if span_streaming:
+        sentry_sdk.flush()
+        spans = [i.payload for i in items if i.type == "span"]
+        assert len(spans) == 1
+        (segment,) = spans
+        assert segment["attributes"]["sentry.origin"] == "auto.http.flask"
+    else:
+        (_, event) = events
+        assert event["contexts"]["trace"]["origin"] == "auto.http.flask"
 
 
-def test_transaction_http_method_default(
+@pytest.mark.parametrize("span_streaming", [True, False])
+def test_transaction_or_segment_http_method_default(
     sentry_init,
     app,
     capture_events,
+    capture_items,
+    span_streaming,
 ):
     """
-    By default OPTIONS and HEAD requests do not create a transaction.
+    By default OPTIONS and HEAD requests do not create a transaction or segment.
     """
     sentry_init(
         traces_sample_rate=1.0,
         integrations=[flask_sentry.FlaskIntegration()],
+        _experiments={"trace_lifecycle": "stream" if span_streaming else "static"},
     )
-    events = capture_events()
+
+    if span_streaming:
+        items = capture_items("span")
+    else:
+        events = capture_events()
 
     client = app.test_client()
     response = client.get("/nomessage")
@@ -1034,16 +1129,25 @@ def test_transaction_http_method_default(
     response = client.head("/nomessage")
     assert response.status_code == 200
 
-    (event,) = events
-
-    assert len(events) == 1
-    assert event["request"]["method"] == "GET"
+    if span_streaming:
+        sentry_sdk.flush()
+        spans = [i.payload for i in items if i.type == "span"]
+        assert len(spans) == 1
+        (segment,) = spans
+        assert segment["attributes"]["http.request.method"] == "GET"
+    else:
+        (event,) = events
+        assert len(events) == 1
+        assert event["request"]["method"] == "GET"
 
 
-def test_transaction_http_method_custom(
+@pytest.mark.parametrize("span_streaming", [True, False])
+def test_transaction_or_segment_http_method_custom(
     sentry_init,
     app,
     capture_events,
+    capture_items,
+    span_streaming,
 ):
     """
     Configure FlaskIntegration to ONLY capture OPTIONS and HEAD requests.
@@ -1058,8 +1162,13 @@ def test_transaction_http_method_custom(
                 )  # capitalization does not matter
             )  # case does not matter
         ],
+        _experiments={"trace_lifecycle": "stream" if span_streaming else "static"},
     )
-    events = capture_events()
+
+    if span_streaming:
+        items = capture_items("span")
+    else:
+        events = capture_events()
 
     client = app.test_client()
     response = client.get("/nomessage")
@@ -1071,8 +1180,15 @@ def test_transaction_http_method_custom(
     response = client.head("/nomessage")
     assert response.status_code == 200
 
-    assert len(events) == 2
-
-    (event1, event2) = events
-    assert event1["request"]["method"] == "OPTIONS"
-    assert event2["request"]["method"] == "HEAD"
+    if span_streaming:
+        sentry_sdk.flush()
+        spans = [i.payload for i in items if i.type == "span"]
+        assert len(spans) == 2
+        (options_segment, head_segment) = spans
+        assert options_segment["attributes"]["http.request.method"] == "OPTIONS"
+        assert head_segment["attributes"]["http.request.method"] == "HEAD"
+    else:
+        assert len(events) == 2
+        (event1, event2) = events
+        assert event1["request"]["method"] == "OPTIONS"
+        assert event2["request"]["method"] == "HEAD"