Skip to content

Commit 9594a8c

Browse files
alexander-alderman-webbalexander-alderman-webb
authored
test: Respect context manager lifecycles in fake_record_sql_queries (#6295)
In the static trace lifecycle, set `span.start_timestamp` and `span.timestamp` immediately after `__enter__` and `__exit__` are called, respectively. In the streaming trace lifecycle, set both `span._start_timestamp` and `span._end_timestamp` immediately after `__enter__` is called, since the query source is added inside the context manager. Then reset `span._end_timestamp` immediately before `__exit__`, as otherwise the span is not captured.
1 parent 75aaebf commit 9594a8c

2 files changed

Lines changed: 29 additions & 45 deletions

File tree

‎tests/integrations/django/test_db_query_data.py‎

Lines changed: 10 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -346,17 +346,16 @@ def test_no_query_source_if_duration_too_short(sentry_init, client, capture_even
346346

347347
class fake_record_sql_queries: # noqa: N801
348348
def __init__(self, *args, **kwargs):
349-
with record_sql_queries(*args, **kwargs) as span:
350-
self.span = span
351-
352-
self.span.start_timestamp = datetime(2024, 1, 1, microsecond=0)
353-
self.span.timestamp = datetime(2024, 1, 1, microsecond=99999)
349+
self._ctx_mgr = record_sql_queries(*args, **kwargs)
354350

355351
def __enter__(self):
352+
self.span = self._ctx_mgr.__enter__()
353+
self.span.start_timestamp = datetime(2024, 1, 1, microsecond=0)
356354
return self.span
357355

358356
def __exit__(self, type, value, traceback):
359-
pass
357+
self._ctx_mgr.__exit__(type, value, traceback)
358+
self.span.timestamp = datetime(2024, 1, 1, microsecond=99999)
360359

361360
with mock.patch(
362361
"sentry_sdk.integrations.django.record_sql_queries",
@@ -404,17 +403,16 @@ def test_query_source_if_duration_over_threshold(sentry_init, client, capture_ev
404403

405404
class fake_record_sql_queries: # noqa: N801
406405
def __init__(self, *args, **kwargs):
407-
with record_sql_queries(*args, **kwargs) as span:
408-
self.span = span
409-
410-
self.span.start_timestamp = datetime(2024, 1, 1, microsecond=0)
411-
self.span.timestamp = datetime(2024, 1, 1, microsecond=101000)
406+
self._ctx_mgr = record_sql_queries(*args, **kwargs)
412407

413408
def __enter__(self):
409+
self.span = self._ctx_mgr.__enter__()
410+
self.span.start_timestamp = datetime(2024, 1, 1, microsecond=0)
414411
return self.span
415412

416413
def __exit__(self, type, value, traceback):
417-
pass
414+
self._ctx_mgr.__exit__(type, value, traceback)
415+
self.span.timestamp = datetime(2024, 1, 1, microsecond=101000)
418416

419417
with mock.patch(
420418
"sentry_sdk.integrations.django.record_sql_queries",

‎tests/integrations/sqlalchemy/test_sqlalchemy.py‎

Lines changed: 19 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -934,25 +934,19 @@ class Person(Base):
934934

935935
class fake_record_sql_queries: # noqa: N801
936936
def __init__(self, *args, **kwargs):
937-
with record_sql_queries_supporting_streaming(
937+
self._ctx_mgr = record_sql_queries_supporting_streaming(
938938
*args, **kwargs
939-
) as span:
940-
self.span = span
941-
942-
if span_streaming:
943-
self.span._start_timestamp = datetime(2024, 1, 1, microsecond=0)
944-
self.span._end_timestamp = datetime(
945-
2024, 1, 1, microsecond=99999
946-
)
947-
else:
948-
self.span.start_timestamp = datetime(2024, 1, 1, microsecond=0)
949-
self.span.timestamp = datetime(2024, 1, 1, microsecond=99999)
939+
)
950940

951941
def __enter__(self):
942+
self.span = self._ctx_mgr.__enter__()
943+
self.span._start_timestamp = datetime(2024, 1, 1, microsecond=0)
944+
self.span._end_timestamp = datetime(2024, 1, 1, microsecond=99999)
952945
return self.span
953946

954947
def __exit__(self, type, value, traceback):
955-
pass
948+
self.span._end_timestamp = None
949+
self._ctx_mgr.__exit__(type, value, traceback)
956950

957951
with mock.patch(
958952
"sentry_sdk.integrations.sqlalchemy.record_sql_queries_supporting_streaming",
@@ -1000,25 +994,18 @@ class Person(Base):
1000994

1001995
class fake_record_sql_queries: # noqa: N801
1002996
def __init__(self, *args, **kwargs):
1003-
with record_sql_queries_supporting_streaming(
997+
self._ctx_mgr = record_sql_queries_supporting_streaming(
1004998
*args, **kwargs
1005-
) as span:
1006-
self.span = span
1007-
1008-
if span_streaming:
1009-
self.span._start_timestamp = datetime(2024, 1, 1, microsecond=0)
1010-
self.span._end_timestamp = datetime(
1011-
2024, 1, 1, microsecond=99999
1012-
)
1013-
else:
1014-
self.span.start_timestamp = datetime(2024, 1, 1, microsecond=0)
1015-
self.span.timestamp = datetime(2024, 1, 1, microsecond=99999)
999+
)
10161000

10171001
def __enter__(self):
1002+
self.span = self._ctx_mgr.__enter__()
1003+
self.span.start_timestamp = datetime(2024, 1, 1, microsecond=0)
10181004
return self.span
10191005

10201006
def __exit__(self, type, value, traceback):
1021-
pass
1007+
self._ctx_mgr.__exit__(type, value, traceback)
1008+
self.span.timestamp = datetime(2024, 1, 1, microsecond=99999)
10221009

10231010
with mock.patch(
10241011
"sentry_sdk.integrations.sqlalchemy.record_sql_queries_supporting_streaming",
@@ -1159,19 +1146,18 @@ class Person(Base):
11591146

11601147
class fake_record_sql_queries: # noqa: N801
11611148
def __init__(self, *args, **kwargs):
1162-
with record_sql_queries_supporting_streaming(
1149+
self._ctx_mgr = record_sql_queries_supporting_streaming(
11631150
*args, **kwargs
1164-
) as span:
1165-
self.span = span
1166-
1167-
self.span.start_timestamp = datetime(2024, 1, 1, microsecond=0)
1168-
self.span.timestamp = datetime(2024, 1, 1, microsecond=101000)
1151+
)
11691152

11701153
def __enter__(self):
1154+
self.span = self._ctx_mgr.__enter__()
1155+
self.span.start_timestamp = datetime(2024, 1, 1, microsecond=0)
11711156
return self.span
11721157

11731158
def __exit__(self, type, value, traceback):
1174-
pass
1159+
self._ctx_mgr.__exit__(type, value, traceback)
1160+
self.span.timestamp = datetime(2024, 1, 1, microsecond=101000)
11751161

11761162
with mock.patch(
11771163
"sentry_sdk.integrations.sqlalchemy.record_sql_queries_supporting_streaming",

0 commit comments

Comments
 (0)