chemical/prop/test_prop.py

from fastapi.testclient import TestClient
from .api import build_tool, ChemicalPropAPI
from typing import Tuple, Optional, List


class ChemicalPropMock(ChemicalPropAPI):
    def __init__(self) -> None:
        self._endpoint = "https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/"

    def get_name_by_cid(self, cid: str, top_k: Optional[int] = None) -> List[str]:
        ans = ["A", "B", "C", "D", "E"]
        if top_k is None:
            top_k = len(ans)
        return ans[:top_k]

    def get_cid_by_struct(self, smiles: str) -> List[str]:
        return ["123"]

    def get_cid_by_name(self, name: str, name_type: Optional[str] = None) -> List[str]:
        return ["123"]

    def get_prop_by_cid(self, cid: str) -> str:
        return {"works": "well"}


app = build_tool({"debug": True, "chemical_prop_api": ChemicalPropMock()})
client = TestClient(app)


def test_get_name():
    response = client.get("/get_name", params={"cid": 123})
    assert response.status_code == 200
    assert response.json() == {"names": ["A", "B", "C"]}


def test_get_all_names():
    response = client.get("/get_allname", params={"cid": 123})
    assert response.status_code == 200
    assert response.json() == {"names": ["A", "B", "C", "D", "E"]}


def test_get_id_by_struct():
    response = client.get("/get_id_by_struct", params={"smiles": "C1=CC=CC=C1"})
    assert response.status_code == 200
    assert response.json() == {"state": "matched", "content": "123"}


def test_get_id():
    response = client.get("/get_id", params={"name": "benzene"})
    assert response.status_code == 200
    assert response.json() == {
        "state": "precise",
        "content": "123",
    }


def test_get_prop():
    response = client.get("/get_prop", params={"cid": "123"})
    assert response.status_code == 200
    assert response.json() == {"works": "well"}