import os import json from PIL import Image from app import predict, process_image # Mock DETECTOR_AVAILABLE to avoid import errors during test import app app.DETECTOR_AVAILABLE = True app.run_detect = lambda args: None # Mock run_detect def test_cleanup(): # Create a large dummy image large_img_path = "test_large_cleanup.png" Image.new('RGB', (2000, 2000), color='green').save(large_img_path) # Create a dummy output file because predict expects it output_path = "temp_result.json" with open(output_path, 'w') as f: json.dump({"prediction": "real", "confidence": 0.9, "elapsed_time": 0.1}, f) # We need to monkeypatch process_image to track the file it creates original_process_image = app.process_image created_files = [] def tracked_process_image(path): new_path = original_process_image(path) if new_path != path: created_files.append(new_path) return new_path app.process_image = tracked_process_image # Run predict try: print("Running predict...") app.predict(large_img_path, "R50_TF") except Exception as e: print(f"Predict failed: {e}") # Check if cropped file was created assert len(created_files) > 0, "Should have created a cropped file" cropped_path = created_files[0] print(f"Cropped path was: {cropped_path}") # Check if cropped file was deleted assert not os.path.exists(cropped_path), f"Cropped file {cropped_path} should have been deleted" print("Cleanup test passed!") # Cleanup original if os.path.exists(large_img_path): os.remove(large_img_path) if os.path.exists(output_path): os.remove(output_path) if __name__ == "__main__": test_cleanup()