diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py index 81d8c17b51..aeab1e38ff 100644 --- a/sentry_sdk/tracing.py +++ b/sentry_sdk/tracing.py @@ -278,6 +278,14 @@ def __exit__( self.finish() self.deactivate() + async def __aenter__(self) -> Span: + return self.__enter__() + + async def __aexit__( + self, ty: Optional[Any], value: Optional[Any], tb: Optional[Any] + ) -> None: + return self.__exit__(ty, value, tb) + @property def description(self) -> Optional[str]: return self.get_attribute(SentrySpanAttribute.DESCRIPTION) diff --git a/tests/tracing/test_integration_tests.py b/tests/tracing/test_integration_tests.py index df6cf57e29..2b738fcf59 100644 --- a/tests/tracing/test_integration_tests.py +++ b/tests/tracing/test_integration_tests.py @@ -1,5 +1,6 @@ import re import sys +import asyncio from unittest import mock import pytest @@ -280,3 +281,26 @@ def test_good_sysexit_doesnt_fail_transaction( assert "status" not in span.get("tags", {}) assert "status" not in event.get("tags", {}) assert event["contexts"]["trace"]["status"] == "ok" + + +@pytest.mark.asyncio +async def test_async_context_manager(sentry_init, capture_events): + """Test that spans work as async context managers""" + sentry_init(traces_sample_rate=1.0) + events = capture_events() + + async with start_span(name="async_transaction") as transaction: + transaction.set_status(SPANSTATUS.OK) + async with start_span(op="async.task", name="async_operation") as span: + span.set_tag("test", "async") + await asyncio.sleep(0.001) + + assert len(events) == 1 + event = events[0] + assert event["transaction"] == "async_transaction" + assert event["contexts"]["trace"]["status"] == "ok" + assert len(event["spans"]) == 1 + span = event["spans"][0] + assert span["op"] == "async.task" + assert span["description"] == "async_operation" + assert span["tags"]["test"] == "async"