tests/test_assets_commands.py

from __future__ import annotations

from pathlib import Path
import json

import pytest

from adam.assets import AssetRegistry
from adam.commands import CommandValidationError, TrainingCommand


class FakeConfig:
    def __init__(self, values: dict) -> None:
        self.values = values

    def get(self, key: str, default=None):
        return self.values.get(key, default)


def test_asset_registry_prefers_exact_friendly_name(tmp_path: Path) -> None:
    exact = tmp_path / "Mario"
    similar = tmp_path / "Mario 2"
    exact.mkdir()
    similar.mkdir()
    registry = AssetRegistry(tmp_path)
    registry.register(kind="dataset", name="Mario", path=str(exact))
    registry.register(kind="dataset", name="Mario 2", path=str(similar))

    assert [item.name for item in registry.find("dataset", "Mario")] == ["Mario"]


def test_training_command_rejects_uncontrolled_fields() -> None:
    with pytest.raises(CommandValidationError, match="Unsupported command fields"):
        TrainingCommand.from_dict(
            {
                "action": "train",
                "trainer": "ddpm",
                "dataset": "dataset",
                "model_name": "model",
                "epochs": 10,
                "shell_command": "unsafe",
            }
        )


def test_training_command_accepts_only_safe_ddpm_options() -> None:
    command = TrainingCommand.from_dict({
        "action": "train", "trainer": "ddpm", "dataset": "dataset", "model_name": "model", "epochs": 10,
        "training_options": {"resolution": 256, "batch_size": 2, "learning_rate": 0.0001},
    })
    assert command.training_options["resolution"] == 256


def test_flow_discovery_recovers_dataset_from_adam_job_history(tmp_path: Path) -> None:
    flow_root = tmp_path / "Flow"
    dataset = tmp_path / "Dataset"
    model = flow_root / "output_flow_models" / "Model"
    dataset.mkdir()
    (model / "unet").mkdir(parents=True)
    (model / "unet" / "config.json").write_text("{}", encoding="utf-8")
    (model / "flow_model_info.json").write_text(
        '{"model_type":"rectified_flow","name":"Friendly Flow","resolution":128}',
        encoding="utf-8",
    )
    (tmp_path / "data").mkdir()
    (tmp_path / "data" / "jobs.json").write_text(json.dumps({"jobs": [{"plan": {"steps": [{
        "tool_id": "flow_trainer", "arguments": {
            "output_dir": str(model), "dataset_dir": str(dataset),
        },
    }]}}]}), encoding="utf-8")
    registry = AssetRegistry(tmp_path)

    registry.discover(FakeConfig({"tool_folders": {"flow_trainer": str(flow_root)}}))

    flow = registry.find("model", "Friendly Flow", trainer="flow")[0]
    assert flow.dataset_id