import sys import os import unittest from unittest.mock import MagicMock, patch from pathlib import Path import tempfile import shutil from PIL import Image # Add root and core directories to python path ROOT_DIR = Path(__file__).resolve().parent.parent sys.path.insert(0, str(ROOT_DIR)) sys.path.insert(0, str(ROOT_DIR / "core")) # Prevent actual background AI loading by mocking warm_up_ai immediately import main main.warm_up_ai = MagicMock() from fastapi.testclient import TestClient from main import app, models, UPLOAD_DIR, PROCESSED_DIR, RESULT_DIR class TestAPIMocked(unittest.TestCase): def setUp(self): # Override directories in main to point to a temporary test folder self.test_dir = tempfile.TemporaryDirectory() self.temp_path = Path(self.test_dir.name) self.orig_upload = main.UPLOAD_DIR self.orig_processed = main.PROCESSED_DIR self.orig_result = main.RESULT_DIR main.UPLOAD_DIR = self.temp_path / "uploads" main.PROCESSED_DIR = self.temp_path / "processed" main.RESULT_DIR = self.temp_path / "results" for d in [main.UPLOAD_DIR, main.PROCESSED_DIR, main.RESULT_DIR]: d.mkdir(parents=True, exist_ok=True) # Set API state to ready models["ready"] = True models["model"] = MagicMock() models["color_model"] = MagicMock() models["device"] = "cpu" self.client = TestClient(app) def tearDown(self): # Restore original directories main.UPLOAD_DIR = self.orig_upload main.PROCESSED_DIR = self.orig_processed main.RESULT_DIR = self.orig_result self.test_dir.cleanup() def test_status(self): response = self.client.get("/status") self.assertEqual(response.status_code, 200) self.assertEqual(response.json(), {"ai_ready": True}) # Test when models not ready models["ready"] = False response = self.client.get("/status") self.assertEqual(response.json(), {"ai_ready": False}) models["ready"] = True def test_get_settings(self): response = self.client.get("/settings") self.assertEqual(response.status_code, 200) data = response.json() self.assertIn("layout", data) self.assertIn("overlays", data) def test_post_settings_validation(self): # Test valid settings update payload = {"retouch": {"enabled": True, "sensitivity": 3.5}} response = self.client.post("/settings", json=payload) self.assertEqual(response.status_code, 200) self.assertEqual(response.json()["status"], "success") # Test invalid type: expecting dict for retouch, but passing float (Bug 4 regression test) payload_invalid = {"retouch": 3.5} response = self.client.post("/settings", json=payload_invalid) self.assertEqual(response.status_code, 400) self.assertIn("error", response.json()) self.assertIn("must be an object", response.json()["error"]) def test_get_frames(self): response = self.client.get("/frames") self.assertEqual(response.status_code, 200) self.assertIn("frames", response.json()) def test_upload_and_process_mocked(self): # 1. Create a dummy JPEG to upload img = Image.new("RGB", (300, 420), (100, 150, 200)) img_byte_arr = io_bytes = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) img.save(img_byte_arr.name) img_byte_arr.close() with open(img_byte_arr.name, "rb") as f: response = self.client.post( "/upload", files={"file": ("test_portrait.jpg", f, "image/jpeg")} ) os.unlink(img_byte_arr.name) self.assertEqual(response.status_code, 200) upload_data = response.json() self.assertIn("id", upload_data) file_id = upload_data["id"] # Verify the file was saved expected_file = main.UPLOAD_DIR / f"{file_id}.jpg" self.assertTrue(expected_file.exists()) # 2. Mock downstream processing steps with patch("restoration.restore_image") as mock_restore, \ patch("process_images.remove_background") as mock_rmbg, \ patch("main.apply_mlp_pil") as mock_color, \ patch("retouch.retouch_image_pil") as mock_retouch, \ patch("main.generate_layout") as mock_layout: # Setup mocks to return standard PIL Images dummy_img = Image.new("RGB", (300, 420), (200, 200, 200)) dummy_rgba = Image.new("RGBA", (300, 420), (200, 200, 200, 255)) mock_restore.return_value = dummy_img mock_rmbg.return_value = dummy_rgba mock_color.return_value = dummy_rgba mock_retouch.return_value = (dummy_rgba, 5) mock_layout.return_value = Image.new("RGB", (1000, 600)) # Run process route process_payload = { "name": "Test User", "id_number": "98765", "do_restore": True, "do_rmbg": True, "do_color": True, "do_retouch": True, "do_crop": True } response = self.client.post(f"/process/{file_id}", data=process_payload) self.assertEqual(response.status_code, 200) process_data = response.json() self.assertIn("result_url", process_data) self.assertIn("preview_url", process_data) # Verify layout was saved layout_file = main.RESULT_DIR / f"{file_id}_layout.jpg" self.assertTrue(layout_file.exists()) if __name__ == "__main__": unittest.main()