test_server.py

import unittest
import numpy as np
import cv2
import torch
from fastapi.testclient import TestClient
from server import app, unet_eye, cls_eye, nail_model
import io
from PIL import Image

class TestAnemiaAPI(unittest.TestCase):
    def setUp(self):
        self.client = TestClient(app)
        
    def create_dummy_image(self, size=(256, 256), color=(255, 0, 0)):
        img = Image.new("RGB", size, color)
        buf = io.BytesIO()
        img.save(buf, format="JPEG")
        return buf.getvalue()

    def test_health_check(self):
        response = self.client.get("/health")
        self.assertEqual(response.status_code, 200)
        self.assertEqual(response.json()["status"], "online")

    def test_predict_endpoint(self):
        # Create dummy images for eye and nail
        eye_img_bytes = self.create_dummy_image()
        nail_img_bytes = self.create_dummy_image()
        
        payload = {
            "gemini_key": "MOCK_KEY"
        }
        
        files = {
            "eye_image": ("eye.jpg", eye_img_bytes, "image/jpeg"),
            "nail_image": ("nail.jpg", nail_img_bytes, "image/jpeg")
        }
        
        response = self.client.post("/predict", data=payload, files=files)
        self.assertEqual(response.status_code, 200)
        
        data = response.json()
        self.assertEqual(data["status"], "success")
        self.assertIn("eye", data)
        self.assertIn("nail", data)
        self.assertIn("anemic_probability", data["eye"])
        self.assertIn("anemic_probability", data["nail"])


    def test_models_loaded(self):
        # Verify eye models are loaded and on correct device
        self.assertIsInstance(unet_eye, torch.nn.Module)
        self.assertIsInstance(cls_eye, torch.nn.Module)
        self.assertIsInstance(nail_model, torch.nn.Module)
        
        # Basic check that models are on the configured device (cuda/cpu)
        # Note: can't easily check actual device unless we inspect params
        pass

if __name__ == "__main__":
    unittest.main()