Add session tracking to ASGI integration (getsentry#1329)

antonpirker · web-flow · commit 435e8567bcce · 2022-02-09T12:58:02.000Z
* test(wsgi): Test for correct session aggregates in wsgi

* test(asgi): added failing test

* feat(asgi): auto session tracking
diff --git a/sentry_sdk/integrations/asgi.py b/sentry_sdk/integrations/asgi.py
@@ -12,6 +12,7 @@
 from sentry_sdk._types import MYPY
 from sentry_sdk.hub import Hub, _should_send_default_pii
 from sentry_sdk.integrations._wsgi_common import _filter_headers
+from sentry_sdk.sessions import auto_session_tracking
 from sentry_sdk.utils import (
     ContextVar,
     event_from_exception,
@@ -119,37 +120,38 @@ async def _run_app(self, scope, callback):
         _asgi_middleware_applied.set(True)
         try:
             hub = Hub(Hub.current)
-            with hub:
-                with hub.configure_scope() as sentry_scope:
-                    sentry_scope.clear_breadcrumbs()
-                    sentry_scope._name = "asgi"
-                    processor = partial(self.event_processor, asgi_scope=scope)
-                    sentry_scope.add_event_processor(processor)
-
-                ty = scope["type"]
-
-                if ty in ("http", "websocket"):
-                    transaction = Transaction.continue_from_headers(
-                        self._get_headers(scope),
-                        op="{}.server".format(ty),
-                    )
-                else:
-                    transaction = Transaction(op="asgi.server")
-
-                transaction.name = _DEFAULT_TRANSACTION_NAME
-                transaction.set_tag("asgi.type", ty)
-
-                with hub.start_transaction(
-                    transaction, custom_sampling_context={"asgi_scope": scope}
-                ):
-                    # XXX: Would be cool to have correct span status, but we
-                    # would have to wrap send(). That is a bit hard to do with
-                    # the current abstraction over ASGI 2/3.
-                    try:
-                        return await callback()
-                    except Exception as exc:
-                        _capture_exception(hub, exc)
-                        raise exc from None
+            with auto_session_tracking(hub, session_mode="request"):
+                with hub:
+                    with hub.configure_scope() as sentry_scope:
+                        sentry_scope.clear_breadcrumbs()
+                        sentry_scope._name = "asgi"
+                        processor = partial(self.event_processor, asgi_scope=scope)
+                        sentry_scope.add_event_processor(processor)
+
+                    ty = scope["type"]
+
+                    if ty in ("http", "websocket"):
+                        transaction = Transaction.continue_from_headers(
+                            self._get_headers(scope),
+                            op="{}.server".format(ty),
+                        )
+                    else:
+                        transaction = Transaction(op="asgi.server")
+
+                    transaction.name = _DEFAULT_TRANSACTION_NAME
+                    transaction.set_tag("asgi.type", ty)
+
+                    with hub.start_transaction(
+                        transaction, custom_sampling_context={"asgi_scope": scope}
+                    ):
+                        # XXX: Would be cool to have correct span status, but we
+                        # would have to wrap send(). That is a bit hard to do with
+                        # the current abstraction over ASGI 2/3.
+                        try:
+                            return await callback()
+                        except Exception as exc:
+                            _capture_exception(hub, exc)
+                            raise exc from None
         finally:
             _asgi_middleware_applied.set(False)
 
diff --git a/tests/integrations/asgi/test_asgi.py b/tests/integrations/asgi/test_asgi.py
@@ -1,7 +1,9 @@
+from collections import Counter
 import sys
 
 import pytest
 from sentry_sdk import Hub, capture_message, last_event_id
+import sentry_sdk
 from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
 from starlette.applications import Starlette
 from starlette.responses import PlainTextResponse
@@ -39,7 +41,7 @@ def test_sync_request_data(sentry_init, app, capture_events):
     events = capture_events()
 
     client = TestClient(app)
-    response = client.get("/sync-message?foo=bar", headers={"Foo": u"ä"})
+    response = client.get("/sync-message?foo=bar", headers={"Foo": "ä"})
 
     assert response.status_code == 200
 
@@ -292,3 +294,41 @@ def test_x_real_ip(sentry_init, app, capture_events):
 
     (event,) = events
     assert event["request"]["env"] == {"REMOTE_ADDR": "1.2.3.4"}
+
+
+def test_auto_session_tracking_with_aggregates(app, sentry_init, capture_envelopes):
+    """
+    Test for correct session aggregates in auto session tracking.
+    """
+
+    @app.route("/dogs/are/great/")
+    @app.route("/trigger/an/error/")
+    def great_dogs_handler(request):
+        if request["path"] != "/dogs/are/great/":
+            1 / 0
+        return PlainTextResponse("dogs are great")
+
+    sentry_init(traces_sample_rate=1.0)
+    envelopes = capture_envelopes()
+
+    app = SentryAsgiMiddleware(app)
+    client = TestClient(app, raise_server_exceptions=False)
+    client.get("/dogs/are/great/")
+    client.get("/dogs/are/great/")
+    client.get("/trigger/an/error/")
+
+    sentry_sdk.flush()
+
+    count_item_types = Counter()
+    for envelope in envelopes:
+        count_item_types[envelope.items[0].type] += 1
+
+    assert count_item_types["transaction"] == 3
+    assert count_item_types["event"] == 1
+    assert count_item_types["sessions"] == 1
+    assert len(envelopes) == 5
+
+    session_aggregates = envelopes[-1].items[0].payload.json["aggregates"]
+    assert session_aggregates[0]["exited"] == 2
+    assert session_aggregates[0]["crashed"] == 1
+    assert len(session_aggregates) == 1
diff --git a/tests/integrations/wsgi/test_wsgi.py b/tests/integrations/wsgi/test_wsgi.py
@@ -3,6 +3,7 @@
 
 import sentry_sdk
 from sentry_sdk.integrations.wsgi import SentryWsgiMiddleware
+from collections import Counter
 
 try:
     from unittest import mock  # python 3.3 and above
@@ -219,7 +220,6 @@ def app(environ, start_response):
 
     traces_sampler = mock.Mock(return_value=True)
     sentry_init(send_default_pii=True, traces_sampler=traces_sampler)
-
     app = SentryWsgiMiddleware(app)
     envelopes = capture_envelopes()
 
@@ -236,3 +236,46 @@ def app(environ, start_response):
     aggregates = sess_event["aggregates"]
     assert len(aggregates) == 1
     assert aggregates[0]["exited"] == 1
+
+
+def test_auto_session_tracking_with_aggregates(sentry_init, capture_envelopes):
+    """
+    Test for correct session aggregates in auto session tracking.
+    """
+
+    def sample_app(environ, start_response):
+        if environ["REQUEST_URI"] != "/dogs/are/great/":
+            1 / 0
+
+        start_response("200 OK", [])
+        return ["Go get the ball! Good dog!"]
+
+    traces_sampler = mock.Mock(return_value=True)
+    sentry_init(send_default_pii=True, traces_sampler=traces_sampler)
+    app = SentryWsgiMiddleware(sample_app)
+    envelopes = capture_envelopes()
+    assert len(envelopes) == 0
+
+    client = Client(app)
+    client.get("/dogs/are/great/")
+    client.get("/dogs/are/great/")
+    try:
+        client.get("/trigger/an/error/")
+    except ZeroDivisionError:
+        pass
+
+    sentry_sdk.flush()
+
+    count_item_types = Counter()
+    for envelope in envelopes:
+        count_item_types[envelope.items[0].type] += 1
+
+    assert count_item_types["transaction"] == 3
+    assert count_item_types["event"] == 1
+    assert count_item_types["sessions"] == 1
+    assert len(envelopes) == 5
+
+    session_aggregates = envelopes[-1].items[0].payload.json["aggregates"]
+    assert session_aggregates[0]["exited"] == 2
+    assert session_aggregates[0]["crashed"] == 1
+    assert len(session_aggregates) == 1
