feat: add cleanup utility for test artifacts in checkpoint resumption tests
- Introduced a new `_cleanup_test_artifacts` function to remove test-generated files and directories after test execution. - Integrated the cleanup function into the `test_current_implementation` and `test_interruption_and_resume` tests to ensure proper resource management and prevent clutter from leftover files.
This commit is contained in:
@@ -105,6 +105,16 @@ def monitor_checkpoint_during_run(checkpoint_file: Path, duration: int = 30) ->
|
||||
return snapshots
|
||||
|
||||
|
||||
def _cleanup_test_artifacts(*paths):
|
||||
"""Remove test-generated files and directories."""
|
||||
for p in paths:
|
||||
p = Path(p)
|
||||
if p.is_dir():
|
||||
shutil.rmtree(p, ignore_errors=True)
|
||||
elif p.is_file():
|
||||
p.unlink(missing_ok=True)
|
||||
|
||||
|
||||
def test_current_implementation():
|
||||
"""Test the current checkpoint implementation."""
|
||||
print("\n" + "=" * 70)
|
||||
@@ -167,6 +177,8 @@ def test_current_implementation():
|
||||
print(f"❌ Error during run: {e}")
|
||||
traceback.print_exc()
|
||||
return False
|
||||
finally:
|
||||
_cleanup_test_artifacts(dataset_file, output_dir)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
@@ -220,9 +232,9 @@ def test_interruption_and_resume():
|
||||
|
||||
print(f"\n▶️ Starting first run (will process 5 prompts, then simulate interruption)...")
|
||||
|
||||
temp_dataset = Path("tests/test_data/checkpoint_test_resume_partial.jsonl")
|
||||
try:
|
||||
# Create a modified dataset with only first 5 prompts for initial run
|
||||
temp_dataset = Path("tests/test_data/checkpoint_test_resume_partial.jsonl")
|
||||
with open(dataset_file, 'r') as f:
|
||||
lines = f.readlines()[:5]
|
||||
with open(temp_dataset, 'w') as f:
|
||||
@@ -292,6 +304,8 @@ def test_interruption_and_resume():
|
||||
print(f"❌ Error during test: {e}")
|
||||
traceback.print_exc()
|
||||
return False
|
||||
finally:
|
||||
_cleanup_test_artifacts(dataset_file, temp_dataset, output_dir)
|
||||
|
||||
|
||||
def test_simulated_crash():
|
||||
|
||||
Reference in New Issue
Block a user