from __future__ import annotations import os import pytest from yolo_webui.config import AugmentationConfig, DatasetSplitConfig, MlflowConfig, TrainingConfig from yolo_webui.trainer import TrainingRunner, mlflow_environment def test_train_kwargs_omit_optional_empty_values() -> None: config = TrainingConfig( dataset=" dataset.yaml ", model=" model.pt ", augmentation=AugmentationConfig(enabled=False), ) from pathlib import Path assert config.train_kwargs() == { "data": "dataset.yaml", "epochs": 100, "imgsz": 640, "batch": 16, "workers": 8, "patience": 100, "project": str(Path("runs/train").resolve()), "verbose": True, } def test_augmentation_kwargs_are_passed_to_ultralytics() -> None: config = TrainingConfig( dataset="dataset.yaml", model="model.pt", augmentation=AugmentationConfig(mosaic=0.5, mixup=0.2, close_mosaic=5), ) kwargs = config.train_kwargs() assert kwargs["mosaic"] == 0.5 assert kwargs["mixup"] == 0.2 assert kwargs["close_mosaic"] == 5 assert kwargs["auto_augment"] == "randaugment" def test_auto_augment_none_policy() -> None: config = TrainingConfig( dataset="dataset.yaml", model="model.pt", augmentation=AugmentationConfig(auto_augment="none"), ) kwargs = config.train_kwargs() assert kwargs["auto_augment"] is None def test_invalid_auto_augment_policy() -> None: config = TrainingConfig( dataset="dataset.yaml", model="model.pt", augmentation=AugmentationConfig(auto_augment="invalid_policy"), # type: ignore[arg-type] ) with pytest.raises(ValueError, match="Неизвестная политика AutoAugment"): config.validate() @pytest.mark.parametrize("field", ["mosaic", "fliplr", "erasing", "perspective"]) def test_augmentation_probabilities_are_validated(field: str) -> None: augmentation = AugmentationConfig(**{field: 1.1}) with pytest.raises(ValueError, match="от 0 до 1"): augmentation.validate() @pytest.mark.parametrize("batch", [0, -2, -0.5, 1.5]) def test_invalid_batch_is_rejected(batch: int | float) -> None: config = TrainingConfig(dataset="dataset.yaml", model="model.pt", batch_size=batch) with pytest.raises(ValueError, match="Batch"): config.validate() @pytest.mark.parametrize("batch", [-1, 1, 16, 0.25]) def test_supported_batch_modes_are_accepted(batch: int | float) -> None: TrainingConfig( dataset="dataset.yaml", model="model.pt", batch_size=batch, ).validate() @pytest.mark.parametrize( ("field", "value", "message"), [ ("epochs", float("nan"), "целым числом"), ("epochs", True, "целым числом"), ("workers", 1.5, "целым числом"), ("patience", float("inf"), "целым числом"), ], ) def test_training_numeric_fields_reject_wrong_types_and_non_finite_values( field: str, value: object, message: str, ) -> None: values = {"dataset": "dataset.yaml", "model": "model.pt", field: value} with pytest.raises(ValueError, match=message): TrainingConfig(**values).validate() def test_augmentation_rejects_non_finite_unbounded_values() -> None: config = TrainingConfig( dataset="dataset.yaml", model="model.pt", augmentation=AugmentationConfig(degrees=float("inf")), ) with pytest.raises(ValueError, match="конечным числом"): config.validate() def test_augmentation_rejects_integer_too_large_for_float() -> None: config = TrainingConfig( dataset="dataset.yaml", model="model.pt", augmentation=AugmentationConfig(degrees=10**10_000), ) with pytest.raises(ValueError, match="конечным числом"): config.validate() @pytest.mark.parametrize( "run_name", ["../escape", "/tmp/escape", r"..\\escape", "C:escape"], ) def test_run_name_cannot_escape_project(run_name: str) -> None: config = TrainingConfig( dataset="dataset.yaml", model="model.pt", run_name=run_name, ) with pytest.raises(ValueError, match="Имя запуска"): config.validate() def test_nested_config_sections_must_be_objects() -> None: with pytest.raises(ValueError, match="разделения датасета.*JSON-объектом"): TrainingConfig.from_dict( { "dataset": "dataset.yaml", "model": "model.pt", "split": None, } ) def test_classes_path_type_is_validated_before_path_operations() -> None: config = TrainingConfig( dataset="dataset.yaml", model="model.pt", split=DatasetSplitConfig(classes_path=123), # type: ignore[arg-type] ) with pytest.raises(ValueError, match="Путь к файлу классов должен быть строкой"): config.validate() def test_mlflow_environment_is_restored(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("MLFLOW_TRACKING_URI", "previous") config = MlflowConfig( enabled=True, tracking_uri="http://127.0.0.1:5000", experiment_name="test-experiment", run_name="test-run", ) with mlflow_environment(config): assert os.environ["MLFLOW_TRACKING_URI"] == config.tracking_uri assert os.environ["MLFLOW_EXPERIMENT_NAME"] == config.experiment_name assert os.environ["MLFLOW_RUN"] == config.run_name assert os.environ["MLFLOW_TRACKING_URI"] == "previous" def test_mlflow_default_tracking_uri_can_be_overridden_by_environment( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("YOLO_WEBUI_MLFLOW_TRACKING_URI", "http://mlflow:5000") assert MlflowConfig().tracking_uri == "http://mlflow:5000" assert "MLFLOW_EXPERIMENT_NAME" not in os.environ assert "MLFLOW_RUN" not in os.environ def test_metrics_summary_skips_non_numeric_values() -> None: summary = TrainingRunner._metrics_summary( { "metrics/mAP50": 0.81234, "label": "invalid", "val/loss": 0.12345, "metrics/precision": 0.95678, "another": 12.34, } ) assert summary == "mAP50=0.8123 · loss=0.1235 · precision=0.9568" @pytest.mark.parametrize("ratio", [0.05, 0.98, float("nan")]) def test_dataset_split_ratio_is_validated(ratio: float) -> None: config = DatasetSplitConfig(enabled=True, train_ratio=ratio) with pytest.raises(ValueError, match="Доля обучающей выборки"): config.validate() def test_classification_rejects_detection_style_auto_split() -> None: config = TrainingConfig( dataset="classification-dataset", model="model.pt", task="classify", split=DatasetSplitConfig(enabled=True), ) with pytest.raises(ValueError, match="classify"): config.validate() @pytest.mark.parametrize( ("field", "value"), [ ("dataset", "https://example.invalid/dataset.yaml"), ("model", "https://example.invalid/model.pt"), ("project", "https://example.invalid/results"), ], ) def test_training_rejects_remote_references(field: str, value: str) -> None: values = { "dataset": "dataset.yaml", "model": "model.pt", "project": "runs/train", field: value, } with pytest.raises(ValueError, match="не может быть URL"): TrainingConfig(**values).validate() def test_model_path_must_stay_in_allowed_roots(tmp_path: Path, monkeypatch) -> None: workspace = tmp_path / "workspace" workspace.mkdir() external_model = tmp_path / "external" / "model.pt" monkeypatch.chdir(workspace) config = TrainingConfig(dataset="dataset.yaml", model=str(external_model)) with pytest.raises(ValueError, match="разрешённом каталоге"): config.validate() monkeypatch.setenv("YOLO_WEBUI_MODEL_ROOTS", str(external_model.parent)) config.validate() def test_existing_bare_dataset_cannot_bypass_allowed_roots( tmp_path: Path, monkeypatch ) -> None: monkeypatch.chdir(tmp_path) (tmp_path / "private.yaml").write_text("secret: value\n", encoding="utf-8") with pytest.raises(ValueError, match="вне разрешённого каталога"): TrainingConfig(dataset="private.yaml", model="model.pt").validate()