tests/test_cli_train_promote.py

"""CLI coverage for checkpoint LoRA promotion."""

from __future__ import annotations

import json
from pathlib import Path

import pytest
from typer.testing import CliRunner

import kneiff.cli.train as cli_train
from kneiff.config import COMFY_MODELS_DIR_ENV_KEY
from tests._cli_helpers import invoke_cli


def test_train_promote_copies_selected_checkpoint_with_release_filename(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
    isolated_cli_project: Path,
) -> None:
    """The command selects a run checkpoint, confirms a target, and preserves source.

    :param monkeypatch: Pytest environment patching helper.
    :param tmp_path: Isolated filesystem root for ComfyUI models.
    :param isolated_cli_project: Fully initialized selected project.
    :return: None.
    """
    workspace = isolated_cli_project / "TRAINING" / "ANIMA_1"
    workspace.mkdir(parents=True)
    (workspace / "simpletuner-config.json").write_text(
        json.dumps({"model_family": "anima", "model_flavour": "base-v1.0"}),
        encoding="utf-8",
    )
    source = (
        workspace
        / "_simpletuner-output"
        / "checkpoint-1400"
        / "pytorch_lora_weights.comfyui.safetensors"
    )
    source.parent.mkdir(parents=True)
    source.write_bytes(b"checkpoint")
    models_dir = tmp_path / "comfy-models"
    destination = models_dir / "models" / "loras" / "characters"
    destination.mkdir(parents=True)
    monkeypatch.setenv(COMFY_MODELS_DIR_ENV_KEY, str(models_dir))
    monkeypatch.setattr(cli_train, "interactive_available", lambda console: True)
    actions = iter(("select", "down", "select", "down", "select", "select"))
    monkeypatch.setattr(cli_train, "read_picker_action", lambda: next(actions))
    monkeypatch.setattr(cli_train.typer, "prompt", lambda *args, **kwargs: "1.2")
    monkeypatch.setattr(cli_train.typer, "confirm", lambda *args, **kwargs: True)

    result = invoke_cli(CliRunner(), ["train", "promote"])

    promoted = destination / "Test_Character-ANIMA-v1.2-1400.safetensors"
    assert result.exit_code == 0, result.output
    assert promoted.read_bytes() == b"checkpoint"
    assert source.read_bytes() == b"checkpoint"
    assert "Target" in result.output
    assert (
        "Promoted LoRA: characters/Test_Character-ANIMA-v1.2-1400.safetensors"
        in result.output
    )


def test_train_promote_name_override_changes_only_filename_identity(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
    isolated_cli_project: Path,
) -> None:
    """The only naming override is applied before target construction.

    :param monkeypatch: Pytest environment patching helper.
    :param tmp_path: Isolated filesystem root for ComfyUI models.
    :param isolated_cli_project: Fully initialized selected project.
    :return: None.
    """
    workspace = isolated_cli_project / "TRAINING" / "F2K_9B_1"
    workspace.mkdir(parents=True)
    (workspace / "simpletuner-config.json").write_text(
        json.dumps({"model_family": "flux2", "model_flavour": "klein-9b"}),
        encoding="utf-8",
    )
    source = (
        workspace / "_simpletuner-output" / "checkpoint-200" / "weights.safetensors"
    )
    source.parent.mkdir(parents=True)
    source.write_bytes(b"checkpoint")
    models_dir = tmp_path / "comfy-models"
    lora_root = models_dir / "models" / "loras"
    lora_root.mkdir(parents=True)
    monkeypatch.setenv(COMFY_MODELS_DIR_ENV_KEY, str(models_dir))
    monkeypatch.setattr(cli_train, "interactive_available", lambda console: True)
    actions = iter(("select", "down", "select", "select"))
    monkeypatch.setattr(cli_train, "read_picker_action", lambda: next(actions))
    monkeypatch.setattr(cli_train.typer, "prompt", lambda *args, **kwargs: "v3.0")
    monkeypatch.setattr(cli_train.typer, "confirm", lambda *args, **kwargs: True)

    result = invoke_cli(CliRunner(), ["train", "promote", "--name", "Rook-Preview"])

    assert result.exit_code == 0, result.output
    assert (lora_root / "Rook-Preview-F2K_9B-v3.0-200.safetensors").is_file()