AlexWelcing/glim-mlip-bench

1"""Smoke tests for hf_space/app.py callable functions (no GPU, no install).2 3These tests stub the elastic-constant calculation so they run without ASE,4torch, CHGNet, etc. They cover the data-shaping logic — payload parsing in5predict_batch, error path in predict, BenchmarkRecord conversion — which is6where bugs would actually hide.7"""8from __future__ import annotations9 10import json11import sys12from pathlib import Path13 14import pytest15 16# Stub the spaces module BEFORE importing app, so the @spaces.GPU decorator17# is a no-op. (This also covers the local-dev path where spaces isn't installed.)18sys.modules.setdefault("spaces", type(sys)("spaces"))19 20# Add hf_space to path21HF_SPACE_DIR = Path(__file__).resolve().parent22sys.path.insert(0, str(HF_SPACE_DIR))23 24 25def test_default_chgnet_is_installed_at_image_build_time() -> None:26    requirements = (HF_SPACE_DIR / "requirements.txt").read_text(encoding="utf-8").splitlines()27    assert "chgnet==0.4.0" in requirements28 29 30@pytest.fixture31def app_module(monkeypatch):32    """Import app.py with elastic.elastic_constants stubbed to a known value."""33    import importlib34 35    class FakeContext:36        def __enter__(self):37            return self38 39        def __exit__(self, *_args):40            return False41 42    class FakeButton:43        def click(self, **_kwargs):44            return None45 46    fake_gradio = type(sys)("gradio")47    fake_gradio.Blocks = lambda *_args, **_kwargs: FakeContext()48    fake_gradio.Tabs = lambda *_args, **_kwargs: FakeContext()49    fake_gradio.TabItem = lambda *_args, **_kwargs: FakeContext()50    fake_gradio.Markdown = lambda *_args, **_kwargs: None51    fake_gradio.Dropdown = lambda *_args, **_kwargs: object()52    fake_gradio.Textbox = lambda *_args, **_kwargs: object()53    fake_gradio.JSON = lambda *_args, **_kwargs: object()54    fake_gradio.Button = lambda *_args, **_kwargs: FakeButton()55    monkeypatch.setitem(sys.modules, "gradio", fake_gradio)56 57    # Stub elastic_constants and ElasticResult before app imports them.58    fake_elastic = type(sys)("elastic")59    class FakeResult:60        def __init__(self, element):61            self.element = element62            self.structure = "fcc" if element in ("Al", "Cu") else "bcc"63            self.a0 = 4.05 if element == "Al" else 3.6264            self.c11 = 108.0 if element == "Al" else 168.065            self.c12 = 61.0 if element == "Al" else 121.066            self.c44 = 28.5 if element == "Al" else 75.067            self.energy_per_atom = -3.36 if element == "Al" else -3.5068    fake_elastic.ElasticResult = FakeResult69    def fake_elastic_constants(element, calc, **_kw):70        if element == "BAD":71            raise RuntimeError("synthetic failure")72        return FakeResult(element)73    fake_elastic.elastic_constants = fake_elastic_constants74    monkeypatch.setitem(sys.modules, "elastic", fake_elastic)75 76    # Stub calculators.make_calculator so it never tries to load real weights.77    fake_calculators = type(sys)("calculators")78    fake_calculators.make_calculator = lambda mlip_id: f"FAKE-{mlip_id}"79    monkeypatch.setitem(sys.modules, "calculators", fake_calculators)80 81    # Force a clean import82    if "app" in sys.modules:83        del sys.modules["app"]84    import app  # noqa: WPS43385    importlib.reload(app)86    app._CALC_CACHE.clear()87    yield app88 89 90def test_predict_single_returns_expected_keys(app_module) -> None:91    out = app_module.predict("Al", "chgnet")92    assert out == {93        "element": "Al", "structure": "fcc",94        "a0": 4.05, "c11": 108.0, "c12": 61.0, "c44": 28.5,95    }96 97 98def test_predict_error_path(app_module) -> None:99    out = app_module.predict("BAD", "chgnet")100    assert "error" in out101    assert out["element"] == "BAD"102    assert out["mlip"] == "chgnet"103 104 105def test_predict_batch_without_refs_returns_simple_records(app_module) -> None:106    out = app_module.predict_batch("Al,Cu", "chgnet")107    assert len(out) == 2108    assert {r["element"] for r in out} == {"Al", "Cu"}109    for r in out:110        assert "c11" in r and "c12" in r and "c44" in r and "a0" in r111        assert "predicted" not in r  # no refs → not in BenchmarkRecord shape112 113 114def test_predict_batch_with_refs_returns_benchmark_records(app_module) -> None:115    refs = json.dumps({116        "Al": {"C11": 108.2, "C12": 61.3, "C44": 28.5, "a0": 4.05},117        "Cu": {"C11": 168.4},118    })119    out = app_module.predict_batch("Al,Cu", "chgnet", refs)120    # Al has 4 props, Cu has 1 → 5 records121    assert len(out) == 5122    al_props = sorted(r["property"] for r in out if r["element"] == "Al")123    assert al_props == ["C11", "C12", "C44", "a0"]124    cu_props = sorted(r["property"] for r in out if r["element"] == "Cu")125    assert cu_props == ["C11"]126    sample = out[0]127    # Must match a BenchmarkRecord-compatible schema so /ingest/batch accepts it.128    for key in ("record_id", "element", "potential_id", "potential_label",129                "pair_style", "property", "reference", "predicted", "unit",130                "provenance", "agent_id", "timestamp"):131        assert key in sample, f"BenchmarkRecord missing {key}"132    assert sample["pair_style"] == "mlip"133 134 135def test_predict_batch_continues_past_one_failure(app_module) -> None:136    out = app_module.predict_batch("Al,BAD,Cu", "chgnet")137    assert len(out) == 3138    bad_record = next(r for r in out if r["element"] == "BAD")139    assert "error" in bad_record140 141 142def test_calculator_cached_across_calls(app_module) -> None:143    app_module.predict("Al", "chgnet")144    app_module.predict("Cu", "chgnet")145    assert app_module._CALC_CACHE == {"chgnet": "FAKE-chgnet"}, \146        "MLIP calculator must be cached across requests"147