Skip to content

Commit 7d1bed6

Browse files
committed
make test cleanup conditional to fix performance regression
1 parent b369dfa commit 7d1bed6

File tree

2 files changed

+69
-26
lines changed

2 files changed

+69
-26
lines changed

sdks/python/apache_beam/typehints/native_type_compatibility.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -342,6 +342,7 @@ def convert_to_beam_type(typ):
342342

343343
# Unwrap Python 3.12 `type` aliases (TypeAliasType) to their underlying value.
344344
# This ensures Beam sees the actual aliased type (e.g., tuple[int, ...]).
345+
import sys
345346
if sys.version_info >= (3, 12) and TypeAliasType is not None:
346347
if isinstance(typ, TypeAliasType): # pylint: disable=isinstance-second-argument-not-valid-type
347348
underlying = getattr(typ, '__value__', None)

sdks/python/conftest.py

Lines changed: 68 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,16 @@ def pytest_addoption(parser):
3434
'--test-pipeline-options',
3535
help='Options to use in test pipelines. NOTE: Tests may '
3636
'ignore some or all of these options.')
37+
parser.addoption(
38+
'--enable-test-cleanup',
39+
action='store_true',
40+
default=None,
41+
help='Enable expensive cleanup operations. Auto-enabled in CI.')
42+
parser.addoption(
43+
'--disable-test-cleanup',
44+
action='store_true',
45+
default=False,
46+
help='Disable expensive cleanup operations even in CI.')
3747

3848

3949
# See pytest.ini for main collection rules.
@@ -101,56 +111,88 @@ def configure_beam_rpc_timeouts():
101111
print("Successfully configured Beam RPC timeouts")
102112

103113

114+
def _running_in_ci():
115+
"""Returns True if running in a CI environment."""
116+
return (
117+
os.getenv('GITHUB_ACTIONS') == 'true' or
118+
os.getenv('CI') == 'true' or
119+
os.getenv('CONTINUOUS_INTEGRATION') == 'true'
120+
)
121+
122+
123+
def _should_enable_test_cleanup(config):
124+
"""Returns True if expensive cleanup operations should run."""
125+
if config.getoption('--disable-test-cleanup'):
126+
result = False
127+
reason = "disabled via --disable-test-cleanup"
128+
elif config.getoption('--enable-test-cleanup'):
129+
result = True
130+
reason = "enabled via --enable-test-cleanup"
131+
else:
132+
if _running_in_ci():
133+
result = True
134+
reason = "CI detected"
135+
else:
136+
result = False
137+
reason = "local development"
138+
139+
# Log once per session
140+
if not hasattr(config, '_cleanup_decision_logged'):
141+
print(f"\n[Test Cleanup] Enabled: {result} ({reason})")
142+
config._cleanup_decision_logged = True
143+
144+
return result
145+
146+
104147
@pytest.fixture(autouse=True)
105-
def ensure_clean_state():
148+
def ensure_clean_state(request):
106149
"""
107-
Ensure clean state before each test
108-
to prevent cross-test contamination.
150+
Ensures clean state between tests to prevent contamination.
151+
152+
Expensive operations (sleeps, extra GC) only run in CI or when
153+
explicitly enabled to keep local tests fast.
109154
"""
110155
import gc
111156
import threading
112157
import time
113158

114-
# Force garbage collection to clean up any lingering resources
115-
gc.collect()
159+
enable_cleanup = _should_enable_test_cleanup(request.config)
160+
161+
if enable_cleanup:
162+
gc.collect()
116163

117-
# Log active thread count for debugging
118164
thread_count = threading.active_count()
119-
if thread_count > 50: # Increased threshold since we see 104 threads
165+
if thread_count > 50:
120166
print(f"Warning: {thread_count} active threads detected before test")
121-
122-
# Force a brief pause to let threads settle
123-
time.sleep(0.5)
124-
gc.collect()
167+
if enable_cleanup:
168+
time.sleep(0.5)
169+
gc.collect()
125170

126171
yield
127172

128-
# Enhanced cleanup after test
129173
try:
130-
# Force more aggressive cleanup
131-
gc.collect()
132-
133-
# Brief pause to let any async operations complete
134-
time.sleep(0.1)
135-
136-
# Additional garbage collection
137-
gc.collect()
174+
if enable_cleanup:
175+
gc.collect()
176+
time.sleep(0.1)
177+
gc.collect()
138178
except Exception as e:
139179
print(f"Warning: Cleanup error: {e}")
140180

141181

142182
@pytest.fixture(autouse=True)
143-
def enhance_mock_stability():
144-
"""Enhance mock stability in DinD environment."""
183+
def enhance_mock_stability(request):
184+
"""Improves mock stability in DinD environment."""
145185
import time
146186

147-
# Brief pause before test to ensure clean mock state
148-
time.sleep(0.05)
187+
enable_cleanup = _should_enable_test_cleanup(request.config)
188+
189+
if enable_cleanup:
190+
time.sleep(0.05)
149191

150192
yield
151193

152-
# Brief pause after test to let mocks clean up
153-
time.sleep(0.05)
194+
if enable_cleanup:
195+
time.sleep(0.05)
154196

155197

156198
def pytest_configure(config):

0 commit comments

Comments
 (0)