tests/test_folding_api.py

import unittest

from carbon_folding_demo.folding_api import FoldingApiClient


class FakeResponse:
    def __init__(self, payload):
        self.payload = payload

    def raise_for_status(self):
        return None

    def json(self):
        return self.payload


class FakeSession:
    def __init__(self):
        self.posts = []
        self.gets = []
        self.poll_count = 0

    def post(self, url, json, headers, timeout):
        self.posts.append((url, json, headers, timeout))
        return FakeResponse({"job_id": "job-1"})

    def get(self, url, headers, timeout):
        self.gets.append((url, headers, timeout))
        if url.endswith("/tools"):
            return FakeResponse({"tools": [{"id": "esmfold"}]})
        self.poll_count += 1
        if self.poll_count == 1:
            return FakeResponse({"job_id": "job-1", "tool_id": "esmfold", "status": "running"})
        return FakeResponse(
            {
                "job_id": "job-1",
                "tool_id": "esmfold",
                "status": "succeeded",
                "result": {
                    "structures": [
                        {
                            "format": "pdb",
                            "content": "ATOM      1  N   ALA A   1",
                            "confidence": {"mean_plddt": 91.2},
                        }
                    ],
                    "metrics": {"runtime_seconds": 1.0},
                    "warnings": ["ok"],
                },
            }
        )


class FoldingApiClientTest(unittest.TestCase):
    def test_submit_and_wait_normalizes_successful_job(self):
        session = FakeSession()
        client = FoldingApiClient("https://fold.example/api", token="secret", timeout_seconds=5, session=session)

        result = client.submit_and_wait(
            tool_id="esmfold",
            entities=[{"id": "A", "type": "protein", "sequence": "MKT"}],
            options={"seed": 1},
            client_metadata={"source": "test"},
            poll_interval=0,
        )

        self.assertEqual(result.job_id, "job-1")
        self.assertEqual(result.tool_id, "esmfold")
        self.assertEqual(result.primary_structure().format, "pdb")
        self.assertEqual(result.primary_structure().confidence["mean_plddt"], 91.2)
        self.assertEqual(session.posts[0][2]["Authorization"], "Bearer secret")
        self.assertEqual(session.posts[0][1]["client_metadata"]["source"], "test")
        self.assertEqual(session.posts[0][3], 60.0)

    def test_request_timeout_uses_configured_ceiling(self):
        self.assertEqual(FoldingApiClient("https://fold.example/api", timeout_seconds=30).request_timeout, 60.0)
        self.assertEqual(FoldingApiClient("https://fold.example/api", timeout_seconds=900).request_timeout, 300.0)

    def test_list_tools_accepts_wrapped_response(self):
        session = FakeSession()
        client = FoldingApiClient("https://fold.example/api", session=session)
        self.assertEqual(client.list_tools(), [{"id": "esmfold"}])


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