train_utility/tests/test_app.py

from __future__ import annotations

import asyncio
import json
import threading

from fastapi.testclient import TestClient
import pytest

import yolo_webui.app as app_module
from yolo_webui.app import (
    ExportManager,
    LiveState,
    TrainingManager,
    app,
    write_json_atomic,
)


def test_get_config_defaults() -> None:
    client = TestClient(app)
    response = client.get("/api/config/defaults")
    assert response.status_code == 200
    data = response.json()
    assert data["dataset"] == "coco8.yaml"
    assert data["model"] == "yolo11n.pt"
    assert data["augmentation"]["enabled"] is True
    assert data["mlflow"]["enabled"] is True


def test_get_status_idle() -> None:
    client = TestClient(app)
    response = client.get("/api/train/status")
    assert response.status_code == 200
    data = response.json()
    assert data["status"] == "idle"
    assert data["epoch"] == 0
    assert data["total_epochs"] == 0
    assert isinstance(data["logs"], list)


def test_start_training_validation_error() -> None:
    client = TestClient(app)
    # Empty dataset is invalid
    bad_config = {
        "dataset": " ",
        "model": "yolo11n.pt",
        "task": "detect"
    }
    response = client.post("/api/train/start", json=bad_config)
    assert response.status_code == 400
    assert "Укажите путь или имя датасета" in response.json()["detail"]


def test_stop_training_when_idle() -> None:
    client = TestClient(app)
    response = client.post("/api/train/stop")
    assert response.status_code == 200
    assert "Запрос на остановку отправлен" in response.json()["message"]


def test_sessions_flow(monkeypatch, tmp_path) -> None:
    client = TestClient(app)
    # Patch sessions directory to use tmp_path
    monkeypatch.setattr("yolo_webui.app.get_sessions_dir", lambda: tmp_path)

    # 1. Get empty sessions list
    response = client.get("/api/sessions")
    assert response.status_code == 200
    assert response.json() == []

    # 2. Save a session
    config = {
        "dataset": "coco8.yaml",
        "model": "yolo11n.pt",
        "task": "detect"
    }
    response = client.post("/api/sessions/my_session", json=config)
    assert response.status_code == 200
    assert "успешно сохранена" in response.json()["message"]

    # 3. List sessions should contain 'my_session'
    response = client.get("/api/sessions")
    assert response.json() == ["my_session"]

    # 4. Load session
    response = client.get("/api/sessions/my_session")
    assert response.status_code == 200
    assert response.json()["dataset"] == "coco8.yaml"

    # 5. Delete session
    response = client.delete("/api/sessions/my_session")
    assert response.status_code == 200
    assert "удалена" in response.json()["message"]

    # 6. List sessions should be empty again
    response = client.get("/api/sessions")
    assert response.json() == []

    # 7. Loading nonexistent session should return 404
    response = client.get("/api/sessions/nonexistent")
    assert response.status_code == 404


def test_invalid_session_config_is_not_saved(monkeypatch, tmp_path) -> None:
    client = TestClient(app)
    monkeypatch.setattr("yolo_webui.app.get_sessions_dir", lambda: tmp_path)

    response = client.post("/api/sessions/broken", json={"unexpected": True})

    assert response.status_code == 400
    assert "датасета" in response.json()["detail"].lower()
    assert not (tmp_path / "broken.json").exists()


def test_atomic_json_write_cleans_temp_file_after_serialization_error(tmp_path) -> None:
    destination = tmp_path / "session.json"
    destination.write_text('{"old": true}', encoding="utf-8")

    with pytest.raises(TypeError):
        write_json_atomic(destination, {"invalid": object()})

    assert destination.read_text(encoding="utf-8") == '{"old": true}'
    assert list(tmp_path.iterdir()) == [destination]


def test_reserved_session_name_cannot_be_overwritten(monkeypatch, tmp_path) -> None:
    client = TestClient(app)
    monkeypatch.setattr("yolo_webui.app.get_sessions_dir", lambda: tmp_path)

    response = client.post("/api/sessions/last_run", json={"dataset": "data"})

    assert response.status_code == 400
    assert "зарезервировано" in response.json()["detail"]


def test_rejected_busy_start_does_not_overwrite_last_run(
    monkeypatch,
    tmp_path,
) -> None:
    client = TestClient(app)
    previous = {"dataset": "previous.yaml", "model": "previous.pt"}
    last_run_path = tmp_path / "last_run.json"
    last_run_path.write_text(json.dumps(previous), encoding="utf-8")
    monkeypatch.setattr("yolo_webui.app.get_sessions_dir", lambda: tmp_path)

    def reject_start(_config) -> None:
        raise ValueError("Обучение уже выполняется.")

    monkeypatch.setattr(app_module.manager, "start_training", reject_start)

    response = client.post(
        "/api/train/start",
        json={"dataset": "coco8.yaml", "model": "yolo11n.pt"},
    )

    assert response.status_code == 400
    assert json.loads(last_run_path.read_text(encoding="utf-8")) == previous


def test_start_api_rejects_non_finite_numeric_value() -> None:
    client = TestClient(app)

    response = client.post(
        "/api/train/start",
        content='{"dataset":"coco8.yaml","model":"yolo11n.pt","epochs":NaN}',
        headers={"content-type": "application/json"},
    )

    assert response.status_code == 400
    assert "целым числом" in response.json()["detail"]


def test_list_models_uses_path_relative_to_runs(monkeypatch, tmp_path) -> None:
    monkeypatch.chdir(tmp_path)
    direct_model = tmp_path / "runs" / "direct.pt"
    nested_model = tmp_path / "runs" / "detect" / "train" / "weights" / "best.pt"
    direct_model.parent.mkdir()
    nested_model.parent.mkdir(parents=True)
    direct_model.write_bytes(b"")
    nested_model.write_bytes(b"")

    response = TestClient(app).get("/api/models")

    assert response.status_code == 200
    run_names = {
        item["name"] for item in response.json() if item["source"] == "runs"
    }
    assert run_names == {"direct.pt", "detect/train/weights/best.pt"}


def test_export_start_rejects_invalid_config_synchronously() -> None:
    response = TestClient(app).post(
        "/api/export/start",
        json={"model": "missing.pt", "format": "not-a-format"},
    )

    assert response.status_code == 400
    assert "Неподдерживаемый формат" in response.json()["detail"]


def test_export_start_passes_canonical_validated_config(
    monkeypatch,
    tmp_path,
) -> None:
    monkeypatch.chdir(tmp_path)
    model_path = tmp_path / "models" / "model.pt"
    model_path.parent.mkdir()
    model_path.write_bytes(b"checkpoint")
    received: list[dict[str, object]] = []
    monkeypatch.setattr(app_module.export_manager, "start_export", received.append)

    response = TestClient(app).post(
        "/api/export/start",
        json={
            "model": "model.pt",
            "format": "ONNX",
            "workspace": 2.5,
        },
    )

    assert response.status_code == 200
    assert received == [
        {
            "model": str(model_path),
            "format": "onnx",
            "imgsz": 640,
            "half": False,
            "int8": False,
            "dynamic": False,
            "simplify": False,
            "batch": 1,
            "workspace": 2.5,
        }
    ]


def test_started_event_does_not_deadlock() -> None:
    training_manager = TrainingManager()
    training_manager.state.status = "preparing"
    event = {
        "kind": "started",
        "message": "Обучение началось.",
        "epoch": 0,
        "total_epochs": 3,
    }
    worker = threading.Thread(
        target=training_manager._handle_subprocess_line,
        args=(f"__YOLO_WEBUI_EVENT__:{json.dumps(event)}",),
    )

    worker.start()
    worker.join(timeout=1)

    assert not worker.is_alive()
    assert training_manager.state.status == "training"


def test_background_broadcast_uses_websocket_event_loop() -> None:
    async def scenario() -> None:
        training_manager = TrainingManager()
        server_thread_id = threading.get_ident()

        class FakeWebSocket:
            def __init__(self) -> None:
                self.messages: list[str] = []
                self.send_thread_ids: list[int] = []
                self.sent = asyncio.Event()

            async def send_text(self, payload: str) -> None:
                self.messages.append(payload)
                self.send_thread_ids.append(threading.get_ident())
                self.sent.set()

        websocket = FakeWebSocket()
        training_manager.add_websocket(websocket)  # type: ignore[arg-type]

        worker = threading.Thread(
            target=training_manager.broadcast,
            args=({"type": "status", "status": "training"},),
        )
        worker.start()
        worker.join(timeout=1)
        assert not worker.is_alive()

        await asyncio.wait_for(websocket.sent.wait(), timeout=1)
        assert json.loads(websocket.messages[0])["status"] == "training"
        assert websocket.send_thread_ids == [server_thread_id]

    asyncio.run(scenario())


def test_process_result_distinguishes_success_cancellation_and_failure() -> None:
    succeeded = TrainingManager()
    succeeded._finalize_process_result(0)
    assert succeeded.state.status == "succeeded"

    cancelled = TrainingManager()
    cancelled.state.stop_requested = True
    cancelled._finalize_process_result(0)
    assert cancelled.state.status == "cancelled"

    failed_after_stop = TrainingManager()
    failed_after_stop.state.stop_requested = True
    failed_after_stop._finalize_process_result(1)
    assert failed_after_stop.state.status == "failed"


def test_live_state_snapshot_is_detached() -> None:
    state = LiveState(
        status="training",
        logs=["first"],
        metrics=[{"epoch": 1, "loss": 0.5}],
    )

    snapshot = state.snapshot()
    state.logs.append("second")
    state.metrics[0]["loss"] = 0.25

    assert snapshot["logs"] == ["first"]
    assert snapshot["metrics"] == [{"epoch": 1, "loss": 0.5}]


def test_export_stop_does_not_relabel_completed_failure_as_cancelled() -> None:
    class CompletedProcess:
        returncode = 1

        def poll(self) -> int:
            return self.returncode

        def terminate(self) -> None:
            raise AssertionError("completed process must not be terminated")

    export = ExportManager()
    export.state.status = "exporting"
    export._process = CompletedProcess()  # type: ignore[assignment]

    export.stop_export()
    export._finalize_process_result(1)

    assert export.state.stop_requested is False
    assert export.state.status == "failed"


def test_export_stopped_while_preparing_never_starts_process(monkeypatch) -> None:
    export = ExportManager()
    export.state.status = "preparing"
    export.stop_export()

    def unexpected_popen(*_args, **_kwargs):
        raise AssertionError("subprocess must not start after cancellation")

    monkeypatch.setattr(app_module.subprocess, "Popen", unexpected_popen)

    export._run_subprocess({"model": "unused.pt", "format": "onnx"})

    assert export.state.status == "cancelled"
    assert export._process is None