netsai/model_optimization

1import io2import pytest3from fastapi.testclient import TestClient4from app.main import app5from PIL import Image6 7client = TestClient(app)8 9def test_predict_success():10    # สร้างรูปจำลองขนาด 224x22411    img = Image.new('RGB', (224, 224), color='red')12    img_byte_arr = io.BytesIO()13    img.save(img_byte_arr, format='JPEG')14    img_bytes = img_byte_arr.getvalue()15 16    response = client.post(17        "/predict",18        files={"file": ("test.jpg", img_bytes, "image/jpeg")}19    )20    assert response.status_code == 20021    assert "predicted_class" in response.json()22    assert response.json()["message"] == "Success"23 24def test_predict_invalid_file():25    # ส่งไฟล์ text แทนรูปภาพ26    response = client.post(27        "/predict",28        files={"file": ("test.txt", b"not an image content", "text/plain")}29    )30    assert response.status_code == 400