Spaces:
Sleeping
Sleeping
File size: 2,813 Bytes
8ab4ff2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | 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()
|