Skip to content

Commit ffe1d38

Browse files
committed
test: add more debugging message
1 parent 9aed207 commit ffe1d38

File tree

1 file changed

+17
-3
lines changed

1 file changed

+17
-3
lines changed

tests/unit/session/test_read_gbq_colab.py

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,19 @@ def test_read_gbq_colab_includes_label():
2929
"""Make sure we can tell direct colab usage apart from regular read_gbq usage."""
3030
import bigframes.core.log_adapter as log_adapter
3131

32+
# Track when get_and_reset_api_methods is called
33+
original_get_and_reset = log_adapter.get_and_reset_api_methods
34+
reset_calls = []
35+
36+
def debug_get_and_reset(dry_run=False):
37+
methods = original_get_and_reset(dry_run)
38+
if methods:
39+
reset_calls.append(methods)
40+
print(f"get_and_reset_api_methods called with: {methods}")
41+
return methods
42+
43+
log_adapter.get_and_reset_api_methods = debug_get_and_reset
44+
3245
# Monkey patch to track if decorator is called
3346
original_add = log_adapter.add_api_method
3447
called_methods = []
@@ -48,11 +61,12 @@ def debug_add_api_method(name):
4861

4962
# Ensure call stack is empty before calling the method
5063
log_adapter._call_stack.clear()
51-
print(f"Call stack before call: {log_adapter._call_stack}")
5264

5365
_ = session._read_gbq_colab("SELECT 'read-gbq-colab-test'")
54-
print(f"Call stack after call: {log_adapter._call_stack}")
55-
print(f"Methods tracked by decorator: {called_methods}")
66+
67+
print(f"Total reset calls: {len(reset_calls)}")
68+
for i, methods in enumerate(reset_calls):
69+
print(f"Reset call {i}: {methods}")
5670

5771
configs = session._job_configs # type: ignore
5872

0 commit comments

Comments
 (0)