tests/test_api.py

import os
import time
import unittest

os.environ.setdefault("FOLD_BACKEND", "stub")
os.environ.setdefault("MAX_PROTEIN_AA", "400")

from fastapi.testclient import TestClient  # noqa: E402

from folding_api_service.app import app  # noqa: E402


class FoldingApiTest(unittest.TestCase):
    def setUp(self):
        self.client = TestClient(app)

    def test_health(self):
        response = self.client.get("/health")
        self.assertEqual(response.status_code, 200)
        self.assertTrue(response.json()["ok"])

    def test_tools(self):
        response = self.client.get("/tools")
        self.assertEqual(response.status_code, 200)
        tools = response.json()["tools"]
        self.assertEqual(tools[0]["id"], "esmfold")
        self.assertEqual(tools[0]["max_protein_aa"], 400)

    def test_submit_and_poll_stub_job(self):
        response = self.client.post(
            "/jobs",
            json={
                "tool_id": "esmfold",
                "entities": [{"id": "A", "type": "protein", "sequence": "MLSDEDFKAVFGMTRSAFANLPLWKQQNLKKEKGLF"}],
                "options": {"msa_mode": "none"},
            },
        )
        self.assertEqual(response.status_code, 200)
        job_id = response.json()["job_id"]

        final = None
        for _ in range(20):
            poll = self.client.get(f"/jobs/{job_id}")
            self.assertEqual(poll.status_code, 200)
            final = poll.json()
            if final["status"] in {"succeeded", "failed"}:
                break
            time.sleep(0.05)

        self.assertIsNotNone(final)
        self.assertEqual(final["status"], "succeeded")
        self.assertIn("ATOM", final["result"]["structures"][0]["content"])
        self.assertEqual(final["result"]["structures"][0]["format"], "pdb")

    def test_validation_errors(self):
        cases = [
            {"tool_id": "omegafold", "entities": [{"id": "A", "type": "protein", "sequence": "MKT"}]},
            {"tool_id": "esmfold", "entities": []},
            {"tool_id": "esmfold", "entities": [{"id": "A", "type": "dna", "sequence": "ATG"}]},
            {"tool_id": "esmfold", "entities": [{"id": "A", "type": "protein", "sequence": ""}]},
            {"tool_id": "esmfold", "entities": [{"id": "A", "type": "protein", "sequence": "M1T"}]},
            {"tool_id": "esmfold", "entities": [{"id": "A", "type": "protein", "sequence": "M" * 401}]},
        ]
        for payload in cases:
            with self.subTest(payload=payload):
                response = self.client.post("/jobs", json={**payload, "options": {}})
                self.assertEqual(response.status_code, 400)

    def test_unknown_job(self):
        response = self.client.get("/jobs/missing")
        self.assertEqual(response.status_code, 404)


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