proteoform_analyzer/tests/test_v340_features.py

"""Tests for v3.4.0: Biomni HPC backend fully removed; every step runs locally.

These tests assert that:
  * the ``_hpc.py`` shim is gone and no module under ``core/`` imports ``biomni``
    or calls ``hpc_run_tool`` / ``hpc_get_job_results`` / ``hpc_search_tools``;
  * ``Boltz2Config`` no longer exposes ``use_hpc`` / ``hpc_tool_id`` and its
    Boltz cache default is a user-home path (not an HPC-cluster path);
  * the new local-run config fields exist (antibody RFAntibody dirs, ThermoMPNN);
  * the CLI exposes and wires the new local flags;
  * the local/skip/graft dispatch paths behave correctly on a box with no GPU,
    no local binaries, and no API key (fold -> graft or skip, dock -> skip,
    ddG -> ESM2 tier, antibody -> clean skip).

The suite is designed to pass on a CPU-only sandbox with none of the external
binaries (boltz, thermompnn, rfantibody, boltzgen) installed.
"""
import ast
import os
import glob
import importlib
import inspect
import tempfile

import pytest

import proteoform_analyzer
from proteoform_analyzer.core.config import (
    Boltz2Config, AntibodyConfig, AnalysisConfig,
)

# Root of the installed package (…/proteoform_analyzer)
PKG_ROOT = os.path.dirname(os.path.abspath(proteoform_analyzer.__file__))
STEPS_DIR = os.path.join(PKG_ROOT, "core", "steps")


# ---------------------------------------------------------------------------
# 0. Version
# ---------------------------------------------------------------------------

def test_version_is_340():
    assert proteoform_analyzer.__version__ == "3.4.0"


# ---------------------------------------------------------------------------
# 1. The HPC shim is deleted
# ---------------------------------------------------------------------------

def test_hpc_shim_deleted():
    assert not os.path.exists(os.path.join(STEPS_DIR, "_hpc.py")), \
        "core/steps/_hpc.py must be deleted in v3.4.0"


def test_hpc_shim_not_importable():
    with pytest.raises(ModuleNotFoundError):
        importlib.import_module("proteoform_analyzer.core.steps._hpc")


# ---------------------------------------------------------------------------
# 2. No active HPC / biomni code references anywhere under core/ (source scan)
# ---------------------------------------------------------------------------

# Substrings that would indicate a *live* HPC/biomni dependency. Historical
# changelog wording ("removed in v3.4.0") is allowed, so we scan for the actual
# call/import tokens rather than the word "HPC".
_FORBIDDEN_TOKENS = (
    "hpc_run_tool",
    "hpc_get_job_results",
    "hpc_search_tools",
    "hpc_get_logs",
    "hpc_cancel_job",
    "get_hpc_run_tool",
    "get_hpc_get_job_results",
    "HpcUnavailable",
    "import biomni",
    "from biomni",
    "from ._hpc",
    "import _hpc",
    ".hpc_tool_id",
    "use_hpc",
)


def _py_files_under(root):
    return glob.glob(os.path.join(root, "**", "*.py"), recursive=True)


def test_no_active_hpc_tokens_in_core():
    core_dir = os.path.join(PKG_ROOT, "core")
    offenders = {}
    for fp in _py_files_under(core_dir):
        with open(fp, "r", encoding="utf-8") as fh:
            src = fh.read()
        hits = [tok for tok in _FORBIDDEN_TOKENS if tok in src]
        if hits:
            offenders[os.path.relpath(fp, PKG_ROOT)] = hits
    assert not offenders, f"Active HPC/biomni tokens found: {offenders}"


def test_no_active_hpc_tokens_in_cli_and_gui():
    offenders = {}
    for fp in (os.path.join(PKG_ROOT, "cli.py"), os.path.join(PKG_ROOT, "gui.py")):
        with open(fp, "r", encoding="utf-8") as fh:
            src = fh.read()
        hits = [tok for tok in _FORBIDDEN_TOKENS if tok in src]
        if hits:
            offenders[os.path.basename(fp)] = hits
    assert not offenders, f"Active HPC/biomni tokens found: {offenders}"


def test_biomni_not_imported_by_package_tree():
    """AST-level: no module in the package has a real `import biomni` statement."""
    offenders = []
    for fp in _py_files_under(PKG_ROOT):
        # skip the test suite itself
        if os.sep + "tests" + os.sep in fp:
            continue
        with open(fp, "r", encoding="utf-8") as fh:
            try:
                tree = ast.parse(fh.read())
            except SyntaxError:
                offenders.append((fp, "SYNTAX ERROR"))
                continue
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for n in node.names:
                    if n.name.split(".")[0] == "biomni":
                        offenders.append((os.path.relpath(fp, PKG_ROOT), n.name))
            elif isinstance(node, ast.ImportFrom):
                if (node.module or "").split(".")[0] == "biomni":
                    offenders.append((os.path.relpath(fp, PKG_ROOT), node.module))
    assert not offenders, f"biomni imports found: {offenders}"


# ---------------------------------------------------------------------------
# 3. Boltz2Config: HPC fields removed, cache default changed
# ---------------------------------------------------------------------------

def test_boltz2config_hpc_fields_removed():
    b = Boltz2Config()
    assert not hasattr(b, "use_hpc"), "Boltz2Config.use_hpc must be removed"
    assert not hasattr(b, "hpc_tool_id"), "Boltz2Config.hpc_tool_id must be removed"


def test_boltz2config_rejects_hpc_kwargs():
    with pytest.raises(TypeError):
        Boltz2Config(use_hpc=True)
    with pytest.raises(TypeError):
        Boltz2Config(hpc_tool_id="boltz-2")


def test_boltz2config_cache_default_is_home():
    b = Boltz2Config()
    # Must no longer point at the old HPC-cluster path.
    assert "/mnt/fsx" not in b.cache_dir
    assert b.cache_dir == os.path.expanduser("~/.cache/boltz")


# ---------------------------------------------------------------------------
# 4. New local-run config fields
# ---------------------------------------------------------------------------

def test_antibody_local_fields_exist():
    a = AntibodyConfig()
    for f in ("local_rfantibody_dir", "local_weights_dir",
              "local_python", "local_framework_pdb"):
        assert hasattr(a, f), f"AntibodyConfig.{f} missing"
        assert getattr(a, f) is None  # default: unconfigured


def test_thermompnn_config_fields_exist():
    c = AnalysisConfig()
    for f in ("thermompnn_dir", "thermompnn_script",
              "thermompnn_checkpoint", "thermompnn_python"):
        assert hasattr(c, f), f"AnalysisConfig.{f} missing"
        assert getattr(c, f) is None


# ---------------------------------------------------------------------------
# 5. CLI exposes + wires the new local flags
# ---------------------------------------------------------------------------

def test_cli_has_new_local_flags():
    from proteoform_analyzer.cli import build_parser
    parser = build_parser()
    help_txt = parser.format_help() if hasattr(parser, "format_help") else ""
    # Parse a representative command line and confirm the fields land on args.
    args = parser.parse_args([
        "run", "--uniprot", "P69905", "--n-subunits", "1",
        "--antibody",
        "--antibody-rfantibody-dir", "/tmp/rfab",
        "--antibody-weights-dir", "/tmp/w",
        "--thermompnn-dir", "/tmp/tmpnn",
        "--thermompnn-checkpoint", "/tmp/ckpt.pt",
    ])
    assert args.antibody_rfantibody_dir == "/tmp/rfab"
    assert args.antibody_weights_dir == "/tmp/w"
    assert args.thermompnn_dir == "/tmp/tmpnn"
    assert args.thermompnn_checkpoint == "/tmp/ckpt.pt"


def test_cli_structure_source_help_has_no_hpc():
    from proteoform_analyzer.cli import build_parser
    parser = build_parser()
    txt = parser.format_help()
    # The word "HPC" must not survive in the run-subcommand help text.
    # (We check the subparser help by formatting the full parser tree.)
    assert "HPC" not in txt


# ---------------------------------------------------------------------------
# 6. Folding dispatch: graft / skip on a bare box
# ---------------------------------------------------------------------------

def _bare_cfg(monkeypatch, **boltz_kw):
    """AnalysisConfig with a Boltz2Config; helper for dispatch tests."""
    ac = AnalysisConfig(uniprot_ids=["P69905"], n_subunits=1)
    ac.boltz2 = Boltz2Config(**boltz_kw)
    return ac


def test_fold_dispatch_functions_present():
    from proteoform_analyzer.core.steps import boltz2_fold as bf
    # Local + graft dispatch helpers must exist; HPC helpers must be gone.
    for fn in ("_fold_local", "_fold_graft", "_fold_api", "build_structures",
               "_resolve_local_binary", "_run_local_boltz"):
        assert hasattr(bf, fn), f"missing {fn}"
    for gone in ("_submit_hpc", "_hpc_command", "_try_collect_job",
                 "_manifest_path", "_load_manifest", "_save_manifest"):
        assert not hasattr(bf, gone), f"{gone} should have been deleted"


def test_build_structures_skips_cleanly_without_backend(monkeypatch):
    """No API key, no local binary, graft disabled -> clean skip (not crash)."""
    from proteoform_analyzer.core.steps import boltz2_fold as bf
    from proteoform_analyzer.core.steps import _boltz_backend as bb

    monkeypatch.setattr(bb, "resolve_backend", lambda cfg, kind: "none")

    cfg = _bare_cfg(monkeypatch, allow_graft_fallback=False)
    with tempfile.TemporaryDirectory() as tmp:
        paths = {k: os.path.join(tmp, k) for k in ("pdbs", "pdbs_monomer")}
        for d in paths.values():
            os.makedirs(d, exist_ok=True)
        res = bf.build_structures(cfg, paths)
    assert res.status in ("skipped", "ok")
    # On a bare box with backend 'none', it must not be 'failed'.
    assert res.status != "failed", res.message


# ---------------------------------------------------------------------------
# 7. Docking: clean skip when backend resolves to none
# ---------------------------------------------------------------------------

def test_docking_boltz2_skips_cleanly_without_backend(monkeypatch):
    from proteoform_analyzer.core.steps import docking as dk
    from proteoform_analyzer.core.steps import _boltz_backend as bb

    # Ensure the boltz2 docking path resolves to 'none' and no local binary.
    monkeypatch.setattr(bb, "resolve_backend", lambda cfg, kind: "none")

    # _run_boltz2_dock is the internal entry; assert it exists and skips.
    assert hasattr(dk, "_run_boltz2_dock")
    src = inspect.getsource(dk)
    # The HPC submit branch must be gone from the docking source.
    for tok in ("hpc_run_tool", "get_hpc_run_tool", "hpc_tool_id"):
        assert tok not in src, f"docking.py still references {tok}"


# ---------------------------------------------------------------------------
# 8. ddG: ThermoMPNN local resolver + ESM2 fallback structure
# ---------------------------------------------------------------------------

def test_ddg_thermompnn_resolver_returns_none_when_unconfigured():
    from proteoform_analyzer.core.steps import ddg
    cfg = AnalysisConfig()  # thermompnn_* all None
    assert ddg._resolve_thermompnn(cfg) is None


def test_ddg_thermompnn_resolver_needs_checkpoint(tmp_path):
    """A dir with a script but no checkpoint must not resolve (falls back to ESM2)."""
    from proteoform_analyzer.core.steps import ddg
    tdir = tmp_path / "thermompnn"
    tdir.mkdir()
    (tdir / "custom_inference.py").write_text("# stub\n")
    cfg = AnalysisConfig(thermompnn_dir=str(tdir))  # no checkpoint
    assert ddg._resolve_thermompnn(cfg) is None


def test_ddg_has_esm2_fallback_and_local_thermompnn():
    from proteoform_analyzer.core.steps import ddg
    for fn in ("_resolve_thermompnn", "_thermompnn_for_chain_local",
               "_run_thermompnn", "_esm2_zeroshot_ddg", "_run_esm2_zeroshot",
               "run_ddg"):
        assert hasattr(ddg, fn), f"missing {fn}"
    # The old HPC ThermoMPNN submit helper must be gone.
    assert not hasattr(ddg, "_thermompnn_for_chain"), \
        "_thermompnn_for_chain (HPC submit) should have been replaced"
    # ddg must not import biomni / call hpc.
    src = inspect.getsource(ddg)
    for tok in ("hpc_run_tool", "biomni", "get_hpc_run_tool"):
        assert tok not in src, f"ddg.py still references {tok}"


# ---------------------------------------------------------------------------
# 9. Antibody: local RFAntibody or clean skip (no HPC)
# ---------------------------------------------------------------------------

def test_antibody_no_hpc_helpers():
    from proteoform_analyzer.core.steps import antibody as ab
    for gone in ("_load_manifest", "_save_manifest", "_manifest_path",
                 "_submit", "_hpc_command"):
        assert not hasattr(ab, gone), f"{gone} should have been deleted"
    for fn in ("_resolve_local_rfantibody", "_stage1_cmd", "_stage2_cmd",
               "_stage3_cmd", "_run_stage", "run_antibody"):
        assert hasattr(ab, fn), f"missing {fn}"
    src = inspect.getsource(ab)
    for tok in ("hpc_run_tool", "biomni", "get_hpc_run_tool", "HpcUnavailable"):
        assert tok not in src, f"antibody.py still references {tok}"


def test_antibody_resolver_none_when_unconfigured():
    from proteoform_analyzer.core.steps.antibody import _resolve_local_rfantibody
    cfg = AnalysisConfig()
    cfg.antibody.enabled = True
    # local_rfantibody_dir is None -> no runnable install
    assert _resolve_local_rfantibody(cfg) is None


def test_antibody_skips_cleanly_without_local_install(tmp_path):
    """Enabled antibody step with a target but no local RFAntibody -> clean skip."""
    from proteoform_analyzer.core.steps.antibody import run_antibody

    paths = {k: str(tmp_path / k) for k in ("antibody", "pdbs", "pdbs_monomer")}
    for d in paths.values():
        os.makedirs(d, exist_ok=True)
    # Minimal WT monomer PDB (two CA atoms, chain A).
    pdb = os.path.join(paths["pdbs_monomer"], "wt-monomer.pdb")
    with open(pdb, "w") as f:
        f.write("ATOM      1  CA  ALA A 305      11.000  11.000  11.000  1.00  0.00           C\n")
        f.write("ATOM      2  CA  GLY A 306      15.000  15.000  15.000  1.00  0.00           C\n")
        f.write("END\n")

    cfg = AnalysisConfig(uniprot_ids=["P69905"], n_subunits=1)
    cfg.antibody.enabled = True
    cfg.antibody.hotspot_source = "user"
    cfg.antibody.hotspot_residues = ["305"]

    res = run_antibody(cfg, paths)
    assert res.status == "skipped"
    assert "local RFAntibody" in res.message
    # No design_context.json is written when we skip before provenance.
    assert not os.path.exists(os.path.join(paths["antibody"], "design_context.json"))


def test_antibody_disabled_skips():
    from proteoform_analyzer.core.steps.antibody import run_antibody
    cfg = AnalysisConfig()
    with tempfile.TemporaryDirectory() as tmp:
        res = run_antibody(cfg, {"antibody": os.path.join(tmp, "ab")})
    assert res.status == "skipped"
    assert "not enabled" in res.message