63 lines
1.8 KiB
Python
Executable file
63 lines
1.8 KiB
Python
Executable file
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from yolo_webui import subprocess_runner
|
|
|
|
|
|
class FakeRunner:
|
|
def __init__(self, *, error: Exception | None = None) -> None:
|
|
self.error = error
|
|
self.prepared = False
|
|
self.stop_requested = False
|
|
|
|
def prepare_run(self) -> None:
|
|
self.prepared = True
|
|
|
|
def request_stop(self) -> None:
|
|
self.stop_requested = True
|
|
|
|
def train(self, config: Any, on_event: Any) -> Path:
|
|
if self.error is not None:
|
|
raise self.error
|
|
return Path("/tmp/successful-run")
|
|
|
|
|
|
def _write_config(tmp_path: Path) -> Path:
|
|
config_path = tmp_path / "config.json"
|
|
config_path.write_text(
|
|
json.dumps({"dataset": "dataset.yaml", "model": "model.pt"}),
|
|
encoding="utf-8",
|
|
)
|
|
return config_path
|
|
|
|
|
|
def test_main_returns_zero_after_successful_training(
|
|
monkeypatch: Any, tmp_path: Path, capsys: Any
|
|
) -> None:
|
|
runner = FakeRunner()
|
|
monkeypatch.setattr(subprocess_runner, "TrainingRunner", lambda: runner)
|
|
|
|
return_code = subprocess_runner.main([str(_write_config(tmp_path))])
|
|
|
|
output = capsys.readouterr()
|
|
assert return_code == 0
|
|
assert runner.prepared is True
|
|
assert "__YOLO_WEBUI_READY__" in output.out
|
|
assert "__YOLO_WEBUI_RESULT__:/tmp/successful-run" in output.out
|
|
assert "Traceback" not in output.err
|
|
|
|
|
|
def test_main_returns_one_when_training_raises(
|
|
monkeypatch: Any, tmp_path: Path, capsys: Any
|
|
) -> None:
|
|
runner = FakeRunner(error=RuntimeError("training failed"))
|
|
monkeypatch.setattr(subprocess_runner, "TrainingRunner", lambda: runner)
|
|
|
|
return_code = subprocess_runner.main([str(_write_config(tmp_path))])
|
|
|
|
output = capsys.readouterr()
|
|
assert return_code == 1
|
|
assert "RuntimeError: training failed" in output.err
|