tests/test_model_plugins.py

from __future__ import annotations

from pathlib import Path
import json

from adam.commands import TrainingCommand
from adam.model_plugins import (
    ModelPluginRegistry,
    scaffold_model_plugin,
    validate_settings,
)
from adam.registry import ToolRegistry


def test_builtin_model_plugins_are_discovered() -> None:
    registry = ModelPluginRegistry(Path.cwd())

    assert {"ddpm", "flow", "lora"}.issubset(registry.plugins)
    assert registry.errors == []
    assert registry.training_schema("ddpm")["resolution"]["type"] == "choice"
    assert registry.generation_schema_for_tool("lora_generator")["base_model_path"]["required"]


def test_plugin_schema_validation_reports_clear_errors(tmp_path: Path) -> None:
    schema = {
        "batch_size": {"label": "Batch size", "type": "int", "min": 1, "max": 8},
        "base_model": {"label": "Base model", "type": "path", "required": True, "must_exist": True},
    }

    errors = validate_settings(schema, {"batch_size": 0, "base_model": str(tmp_path / "missing.safetensors")})

    assert "Batch size must be at least 1." in errors
    assert "Base model must point to an existing file." in errors


def test_plugin_schema_extends_existing_tool_arguments_without_required_breakage() -> None:
    registry = ToolRegistry(Path.cwd())
    lora = registry.get("lora_trainer")

    assert "rank" in lora.arguments
    assert "alpha" in lora.arguments
    assert set(lora.required_arguments) == {
        "dataset_dir",
        "model_name",
        "epochs",
        "output_dir",
        "base_model",
    }


def test_scaffolded_plugin_is_discovered_and_gets_standard_training_arguments(tmp_path: Path) -> None:
    folder = scaffold_model_plugin(
        tmp_path,
        plugin_id="Neural Cellular Automata",
        name="Neural Cellular Automata",
        architecture="nca",
    )

    config = tmp_path / "config"
    config.mkdir()
    (config / "tools.json").write_text(json.dumps({"tools": []}), encoding="utf-8")
    registry = ToolRegistry(tmp_path)
    tool = registry.get("neural_cellular_automata_trainer")

    assert folder.name == "neural_cellular_automata"
    assert "dataset_dir" in tool.arguments
    assert "output_dir" in tool.required_arguments
    assert registry.model_plugins.training_schema("neural_cellular_automata")["resolution"]["default"] == 256


def test_training_command_accepts_discovered_custom_plugin(tmp_path: Path, monkeypatch) -> None:
    scaffold_model_plugin(
        tmp_path,
        plugin_id="maskgit",
        name="MaskGIT",
        architecture="maskgit",
    )
    monkeypatch.chdir(tmp_path)

    command = TrainingCommand.from_dict(
        {
            "action": "train",
            "trainer": "maskgit",
            "dataset": "D:/data",
            "model_name": "Mask Test",
            "epochs": 5,
            "training_options": {"resolution": 256, "batch_size": 1},
        }
    )

    assert command.trainer == "maskgit"