tests/test_api_mocked.py

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()