Blame for tests/test_bankml_push.py

"""push_to_bankml: the Modelfile subset, the capability probe, and `bankml create` — no binary.

subprocess.run and shutil.which are monkeypatched; the LoRA merge is replaced by a stub.
"""

from __future__ import annotations

import json
import subprocess
from pathlib import Path
from typing import Any

import pytest
from typer.testing import CliRunner

from mindxtrain.deploy import bankml_push as BP
from mindxtrain.deploy.modelfile import ModelfileSpec, render_modelfile

SHA = "6b64c748d96ad26fd72402299bd27b2ae82f489bd0469498dd18eb6054058266"
DIGEST = "c" * 64
USAGE_035 = (
    "usage: bankml usage [PID …]\n"
    "       bankml create NAME -f Modelfile [--registry DIR] [--models DIR]\n"
    "       bankml convert SAFETENSORS_DIR -o OUT.gguf [--outtype f16]\n"
    "       bankml version"
)
USAGE_034 = "usage: bankml usage [PID …]\n       bankml serve FILE --fork FORK.json\n       bankml version"


class FakeBankml:
    """Records argv; answers `version`, `--help`, `convert`, `create`, `sha256` like bankml."""

    def __init__(self, usage: str = USAGE_035, create_rc: int = 0, create_err: str = "",
                 version: str = "0.3.5") -> None:
        self.calls: list[list[str]] = []
        self.version = version
        self.usage = usage
        self.create_rc = create_rc
        self.create_err = create_err
        self.modelfile_text = ""

    def __call__(self, cmd: list[str], **_: Any) -> subprocess.CompletedProcess[str]:
        self.calls.append(list(cmd))
        verb = cmd[1]
        if verb == "version":
            return subprocess.CompletedProcess(cmd, 0, f"bankml {self.version}\n", "")
        if verb == "--help":
            return subprocess.CompletedProcess(cmd, 1, "", self.usage)
        if verb == "convert":
            out = cmd[cmd.index("-o") + 1]
            return subprocess.CompletedProcess(cmd, 0, f"{SHA}  {out}\n", "bankml convert: ok\n")
        if verb == "create":
            self.modelfile_text = Path(cmd[cmd.index("-f") + 1]).read_text()
            if self.create_rc:
                return subprocess.CompletedProcess(cmd, self.create_rc, "parsing modelfile\n",
                                                   self.create_err)
            name = cmd[2]
            return subprocess.CompletedProcess(
                cmd, 0, f"parsing modelfile\npinned {name}-F16.gguf sha256:{SHA}\n",
                f"bankml create: {name} (digest sha256:{DIGEST}) over {name}-F16.gguf "
                f"(sha256 {SHA})\n")
        if verb == "sha256":
            return subprocess.CompletedProcess(cmd, 0, f"{SHA}  {cmd[2]}\n", "")
        raise AssertionError(cmd)


@pytest.fixture
def fake(monkeypatch: pytest.MonkeyPatch) -> FakeBankml:
    f = FakeBankml()
    monkeypatch.setattr(BP.shutil, "which", lambda name: "/usr/local/bin/bankml")
    monkeypatch.setattr(BP.subprocess, "run", f)
    return f


def _merged(tmp_path: Path, arch: str = "LlamaForCausalLM") -> Path:
    d = tmp_path / "merged"
    d.mkdir()
    (d / "config.json").write_text(json.dumps({"architectures": [arch]}))
    return d


# ---- sanitize: the refusal table ----------------------------------------------


@pytest.mark.parametrize(
    ("spec_kw", "needle"),
    [
        ({"adapter": "./lora"}, "ADAPTER"),
        ({"template": "{{ .Prompt }}"}, "TEMPLATE"),
        ({"parameters": {"repeat_penalty": 1.3}}, "PARAMETER repeat_penalty: penalties"),
        ({"parameters": {"presence_penalty": 0.0}}, "PARAMETER presence_penalty"),
        ({"parameters": {"frequency_penalty": 0.5}}, "PARAMETER frequency_penalty"),
        ({"parameters": {"repeat_last_n": 64}}, "PARAMETER repeat_last_n"),
        ({"parameters": {"mirostat": 2}}, "mirostat is not reproduced"),
        ({"parameters": {"mirostat_tau": 5.0}}, "mirostat is not reproduced"),
        ({"parameters": {"typical_p": 0.9}}, "typical_p"),
        ({"parameters": {"num_thread": 4}}, "resource option"),
        ({"parameters": {"num_gpu": 0}}, "resource option"),
        ({"parameters": {"stop": "x"}}, "ModelfileSpec.stop"),
        ({"parameters": {"weird": 1}}, "not a parameter bankml reproduces"),
        ({"system": 'say """hi"""'}, "SYSTEM"),
    ],
)
def test_sanitize_refuses_and_names_it(spec_kw: dict[str, Any], needle: str) -> None:
    res = BP.bankml_sanitize(ModelfileSpec(from_model="x", **spec_kw))
    assert not res.ok
    assert any(needle in r for r in res.refusals), res.refusals


def test_sanitize_accepts_the_subset_unchanged() -> None:
    spec = ModelfileSpec(
        from_model="/m", system="You are mindX.", license="Apache-2.0",
        parameters={"temperature": 0.0, "top_k": 40, "top_p": 0.9, "min_p": 0.05, "seed": 42,
                    "num_ctx": 2048, "num_predict": 128},
        stop=["<|im_end|>"],
        messages=[{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}],
    )
    res = BP.bankml_sanitize(spec)
    assert res.ok and res.spec == spec
    text = render_modelfile(res.spec)
    assert "PARAMETER seed 42" in text and 'PARAMETER stop "<|im_end|>"' in text


def test_template_passes_only_when_equal_to_base() -> None:
    spec = ModelfileSpec(from_model="x", template="T")
    assert BP.bankml_sanitize(spec, base_template="T").ok
    assert not BP.bankml_sanitize(spec, base_template="U").ok


def test_arch_checks(tmp_path: Path) -> None:
    assert BP.check_merged_arch(_merged(tmp_path)) is None
    q = tmp_path / "q"
    q.mkdir()
    (q / "config.json").write_text(json.dumps({"architectures": ["Qwen3ForCausalLM"]}))
    assert "Llama-architecture" in (BP.check_merged_arch(q) or "")
    assert BP.base_family_refusal("HuggingFaceTB/SmolLM2-135M") is None
    assert "qwen" in (BP.base_family_refusal("Qwen/Qwen3-1.7B") or "")


# ---- capability probe ----------------------------------------------------------


def test_missing_binary(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
    monkeypatch.setattr(BP.shutil, "which", lambda name: None)
    res = BP.push_to_bankml("b", "t", merged_dir=_merged(tmp_path))
    assert res.status == "bankml_missing"
    assert "github.com/cryptoAGI/bankml" in res.reason


def test_too_old_bankml_is_reported_not_raised(fake: FakeBankml, tmp_path: Path) -> None:
    fake.usage = USAGE_034
    res = BP.push_to_bankml("b", "t", merged_dir=_merged(tmp_path))
    assert res.status == "bankml_too_old"
    assert "0.3.5" in res.reason and "create" in res.reason
    assert not any(c[1] == "create" for c in fake.calls)


def test_capabilities_read_from_usage(fake: FakeBankml) -> None:
    caps = BP.bankml_capabilities()
    assert caps.version == "0.3.5" and caps.has_create and caps.has_convert


# ---- push ---------------------------------------------------------------------


def test_push_from_merged_dir(fake: FakeBankml, tmp_path: Path) -> None:
    merged = _merged(tmp_path)
    res = BP.push_to_bankml(
        "HuggingFaceTB/SmolLM2-135M", "mindx-gen99", merged_dir=merged,
        system="You are mindX.", params={"temperature": 0.0, "seed": 1}, stop=["<|im_end|>"],
        registry_dir=tmp_path / "forks", work_dir=tmp_path / "work",
    )
    assert res.ok, res
    assert res.model_sha256 == SHA and res.digest == DIGEST
    create = next(c for c in fake.calls if c[1] == "create")
    assert create[2] == "mindx-gen99" and "--registry" in create
    # converted through a link named after the tag: the GGUF's general.name is "Mindx Gen99"
    link = tmp_path / "work" / "mindx-gen99" / "mindx-gen99"
    assert fake.modelfile_text.startswith(f"FROM {link}\n") and link.resolve() == merged.resolve()
    assert 'SYSTEM """You are mindX."""' in fake.modelfile_text
    assert "repeat_penalty" not in fake.modelfile_text


def test_push_with_convert_first(fake: FakeBankml, tmp_path: Path) -> None:
    res = BP.push_to_bankml("b", "gen7", merged_dir=_merged(tmp_path), convert=True,
                            registry_dir=tmp_path / "forks", work_dir=tmp_path / "work")
    assert res.ok
    conv = next(c for c in fake.calls if c[1] == "convert")
    assert conv[conv.index("-o") + 1].endswith("gen7-base-F16.gguf") and "--fork" in conv
    assert fake.modelfile_text.startswith("FROM ") and "gen7-base-F16.gguf" in fake.modelfile_text


def test_refused_params_never_reach_bankml(fake: FakeBankml, tmp_path: Path) -> None:
    res = BP.push_to_bankml("b", "t", merged_dir=_merged(tmp_path),
                            params={"mirostat": 2, "temperature": 0.7})
    assert res.status == "refused"
    assert any("mirostat" in r for r in res.refusals)
    assert fake.calls == []  # refused before even probing the binary


def test_penalties_refused_before_0_3_6(fake: FakeBankml, tmp_path: Path) -> None:
    res = BP.push_to_bankml("b", "t", merged_dir=_merged(tmp_path),
                            params={"repeat_penalty": 1.3, "temperature": 0.7})
    assert res.status == "refused" and res.bankml_version == "0.3.5"
    assert any("repeat_penalty" in r and "0.3.6" in r for r in res.refusals)
    assert not any(c[1] in ("create", "convert") for c in fake.calls)


def test_penalties_taken_from_0_3_6(fake: FakeBankml, tmp_path: Path) -> None:
    fake.version = "0.3.6"
    res = BP.push_to_bankml("b", "mindx-gen99", merged_dir=_merged(tmp_path), work_dir=tmp_path / "w",
                            params={"repeat_penalty": 1.3, "num_ctx": 2048}, stop=["<|im_end|>"])
    assert res.ok, res
    assert "PARAMETER repeat_penalty 1.3" in fake.modelfile_text


def test_version_gate() -> None:
    assert BP.penalties_supported("0.3.6") and BP.penalties_supported("0.4.0")
    assert not BP.penalties_supported("0.3.5") and not BP.penalties_supported("")
    assert BP.bankml_sanitize(ModelfileSpec(from_model="x", parameters={"repeat_penalty": 1.3})).refusals
    assert BP.bankml_sanitize(ModelfileSpec(from_model="x", parameters={"repeat_penalty": 1.3}),
                              penalties=True).ok


def test_merged_dir_already_named_after_the_tag_is_used_as_is(fake: FakeBankml, tmp_path: Path) -> None:
    merged = tmp_path / "mindx-gen99"
    merged.mkdir()
    (merged / "config.json").write_text(json.dumps({"architectures": ["LlamaForCausalLM"]}))
    res = BP.push_to_bankml("b", "mindx-gen99", merged_dir=merged, work_dir=tmp_path / "w")
    assert res.ok and fake.modelfile_text.startswith(f"FROM {merged.resolve()}\n")


def test_non_llama_merged_dir_refused(fake: FakeBankml, tmp_path: Path) -> None:
    res = BP.push_to_bankml("b", "t", merged_dir=_merged(tmp_path, "Qwen3ForCausalLM"))
    assert res.status == "refused" and "Llama-architecture" in res.reason
    assert not any(c[1] == "create" for c in fake.calls)


def test_create_refusal_is_a_result(fake: FakeBankml, tmp_path: Path) -> None:
    fake.create_rc = 2
    fake.create_err = "bankml create: refuse: FROM /x: tokenizer not seen\n"
    res = BP.push_to_bankml("b", "t", merged_dir=_merged(tmp_path), work_dir=tmp_path / "w")
    assert res.status == "refused"
    assert "tokenizer not seen" in res.reason


def test_bad_tag_and_bad_inputs(tmp_path: Path) -> None:
    assert BP.push_to_bankml("b", "Bad:tag", merged_dir=tmp_path).status == "refused"
    assert BP.push_to_bankml("b", "ok").status == "error"


def test_adapter_merge_and_import_error(fake: FakeBankml, tmp_path: Path,
                                        monkeypatch: pytest.MonkeyPatch) -> None:
    from mindxtrain.deploy import ollama_push

    def merged_stub(base: str, adapter: Path, out: Path, *, sink: Any = None) -> Path:
        out.mkdir(parents=True)
        (out / "config.json").write_text(json.dumps({"architectures": ["LlamaForCausalLM"]}))
        return out

    adapter = tmp_path / "checkpoint"
    adapter.mkdir()
    monkeypatch.setattr(ollama_push, "merge_lora_adapter", merged_stub)
    res = BP.push_to_bankml("HuggingFaceTB/SmolLM2-135M", "g", adapter_dir=adapter,
                            work_dir=tmp_path / "w", registry_dir=tmp_path / "r")
    assert res.ok and res.merged_dir == tmp_path / "w" / "merged"

    def no_ml(*a: Any, **k: Any) -> Path:
        raise ImportError("push-to-ollama needs `peft` + `transformers`")

    monkeypatch.setattr(ollama_push, "merge_lora_adapter", no_ml)
    res = BP.push_to_bankml("b", "g2", adapter_dir=adapter, work_dir=tmp_path / "w2")
    assert res.status == "merge_failed" and "uv sync --extra ml" in res.reason


def test_register_with_mindx_uses_bankml_provider(fake: FakeBankml, tmp_path: Path,
                                                  monkeypatch: pytest.MonkeyPatch) -> None:
    from mindxtrain.deploy import api_client

    seen: dict[str, Any] = {}

    def swap(**kw: Any) -> dict[str, str]:
        seen.update(kw)
        return {"previous": "ollama:x", "current": "bankml:t"}

    monkeypatch.setattr(api_client, "swap_mindx_fallback_model", swap)
    res = BP.push_to_bankml("b", "t", merged_dir=_merged(tmp_path), register_with_mindx=True,
                            work_dir=tmp_path / "w", registry_dir=tmp_path / "r")
    assert res.ok and seen["provider"] == "bankml" and seen["model"] == "t"
    assert res.mindx_fallback_swap == {"previous": "ollama:x", "current": "bankml:t"}


# ---- CLI: serve --to bankml ------------------------------------------------------


def _recipe(tmp_path: Path, **overrides: Any) -> Path:
    import yaml

    from mindxtrain.config.loader import render_recipe

    data = yaml.safe_load(render_recipe("mindx_fallback_qwen3_1_5b_cpu_real"))
    for dotted, value in overrides.items():
        node = data
        *parents, leaf = dotted.split(".")
        for p in parents:
            node = node[p]
        node[leaf] = value
    path = tmp_path / "run.yaml"
    path.write_text(yaml.safe_dump(data))
    return path


def test_cli_refuses_quantized_config(tmp_path: Path) -> None:
    from mindxtrain.cli.main import app

    cfg = _recipe(tmp_path, **{"quantize.enabled": True, "quantize.scheme": "quark_fp8"})
    r = CliRunner().invoke(app, ["serve", str(cfg), "--to", "bankml"])
    assert r.exit_code == 2
    assert "bankml refuses quantize.scheme=quark_fp8" in r.output


def test_cli_refuses_unconvertible_family(tmp_path: Path) -> None:
    from mindxtrain.cli.main import app

    cfg = _recipe(tmp_path, **{"model.name": "Qwen/Qwen3-1.7B"})
    r = CliRunner().invoke(app, ["serve", str(cfg), "--to", "bankml"])
    assert r.exit_code == 2
    assert "qwen" in r.output.lower()


def test_cli_too_old_exits_2(tmp_path: Path, fake: FakeBankml) -> None:
    from mindxtrain.cli.main import app

    fake.usage = USAGE_034
    ckpt = tmp_path / "ck"
    ckpt.mkdir()
    (ckpt / "config.json").write_text(json.dumps({"architectures": ["LlamaForCausalLM"]}))
    cfg = _recipe(tmp_path)
    r = CliRunner().invoke(app, ["serve", str(cfg), "--to", "bankml", "--checkpoint", str(ckpt)])
    assert r.exit_code == 2
    assert "bankml_too_old" in r.output