train_utility/tests/test_app.py

from __future__ import annotations

import asyncio
import json
import threading

from fastapi.testclient import TestClient

from yolo_webui.app import TrainingManager, app


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_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_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"