tests/unit/test_model_trainer.py

"""
Unit tests for ModelTrainer (app/ml/trainer.py).

Uses a small synthetic dataset to verify:
- All four .joblib artifacts are created
- ModelLogRepository.save is called four times (once per algorithm)
- The highest-accuracy model has deployed=True in the DB log
- TrainedModel dataclass fields are populated correctly
- train_all works without a model_log_repo (DB saving skipped)

Requirements: 5.2, 5.5
"""

from __future__ import annotations

import json
import os
import tempfile
from pathlib import Path
from unittest.mock import MagicMock, call, patch

import numpy as np
import pandas as pd
import pytest

from app.ml.trainer import ModelTrainer, TrainedModel, _ARTIFACT_NAMES, _FEATURE_NAMES
from app.ml.evaluator import EvaluationMetrics


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_synthetic_csv(path: str, n_samples: int = 200) -> None:
    """Write a minimal ILPD-format CSV to *path* for testing."""
    rng = np.random.default_rng(42)

    # Build a DataFrame with the original ILPD column names
    n = n_samples
    df = pd.DataFrame(
        {
            "Age": rng.integers(20, 80, size=n),
            "Gender": rng.choice(["Male", "Female"], size=n),
            "Total_Bilirubin": rng.uniform(0.4, 5.0, size=n),
            "Direct_Bilirubin": rng.uniform(0.1, 2.0, size=n),
            "Alkaline_Phosphotase": rng.integers(60, 400, size=n),
            "Alamine_Aminotransferase": rng.integers(10, 200, size=n),
            "Aspartate_Aminotransferase": rng.integers(10, 300, size=n),
            "Total_Protiens": rng.uniform(5.0, 9.0, size=n),
            "Albumin": rng.uniform(2.0, 5.5, size=n),
            "Albumin_and_Globulin_Ratio": rng.uniform(0.5, 2.5, size=n),
            # Target: 1 = liver disease, 2 = no liver disease
            "Dataset": rng.choice([1, 2], size=n),
        }
    )
    df.to_csv(path, index=False)


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

@pytest.fixture()
def tmp_dirs():
    """Provide temporary directories for artifacts, preprocessor, and mask."""
    with tempfile.TemporaryDirectory() as artifacts_dir:
        preprocessor_path = os.path.join(artifacts_dir, "preprocessor_pipeline.joblib")
        mask_path = os.path.join(artifacts_dir, "feature_mask.joblib")
        yield {
            "artifacts_dir": artifacts_dir,
            "preprocessor_path": preprocessor_path,
            "mask_path": mask_path,
        }


@pytest.fixture()
def dataset_csv(tmp_dirs):
    """Write a synthetic CSV and return its path."""
    csv_path = os.path.join(tmp_dirs["artifacts_dir"], "dataset.csv")
    _make_synthetic_csv(csv_path, n_samples=200)
    return csv_path


@pytest.fixture()
def trainer(tmp_dirs):
    """Return a ModelTrainer configured to use temporary directories."""
    return ModelTrainer(
        artifacts_dir=tmp_dirs["artifacts_dir"],
        preprocessor_path=tmp_dirs["preprocessor_path"],
        mask_path=tmp_dirs["mask_path"],
    )


# ---------------------------------------------------------------------------
# Tests: TrainedModel dataclass
# ---------------------------------------------------------------------------

class TestTrainedModelDataclass:
    def test_fields_accessible(self):
        metrics = EvaluationMetrics(
            accuracy=0.8, precision=0.75, recall=0.70,
            f1_score=0.72, tp=40, tn=40, fp=10, fn=10,
        )
        tm = TrainedModel(
            algorithm="random_forest",
            model=object(),
            metrics=metrics,
            cv_score=0.78,
            artifact_path="/tmp/rf.joblib",
        )
        assert tm.algorithm == "random_forest"
        assert tm.cv_score == 0.78
        assert tm.artifact_path == "/tmp/rf.joblib"
        assert tm.metrics.accuracy == 0.8


# ---------------------------------------------------------------------------
# Tests: train_all without DB repo
# ---------------------------------------------------------------------------

class TestTrainAllNoRepo:
    """Verify training works end-to-end without a model_log_repo."""

    def test_returns_dict_with_four_algorithms(self, trainer, dataset_csv):
        results = trainer.train_all(dataset_csv)
        assert set(results.keys()) == {
            "logistic_regression", "decision_tree", "random_forest", "svm"
        }

    def test_each_result_is_trained_model(self, trainer, dataset_csv):
        results = trainer.train_all(dataset_csv)
        for algo, tm in results.items():
            assert isinstance(tm, TrainedModel), f"{algo} result is not a TrainedModel"

    def test_all_four_joblib_artifacts_created(self, trainer, dataset_csv, tmp_dirs):
        trainer.train_all(dataset_csv)
        for algo, filename in _ARTIFACT_NAMES.items():
            artifact_path = os.path.join(tmp_dirs["artifacts_dir"], filename)
            assert os.path.exists(artifact_path), (
                f"Artifact missing for {algo}: {artifact_path}"
            )

    def test_artifact_paths_in_results(self, trainer, dataset_csv, tmp_dirs):
        results = trainer.train_all(dataset_csv)
        for algo, tm in results.items():
            expected = os.path.join(tmp_dirs["artifacts_dir"], _ARTIFACT_NAMES[algo])
            assert tm.artifact_path == expected

    def test_cv_scores_are_floats_in_unit_interval(self, trainer, dataset_csv):
        results = trainer.train_all(dataset_csv)
        for algo, tm in results.items():
            assert isinstance(tm.cv_score, float), f"{algo} cv_score is not float"
            assert 0.0 <= tm.cv_score <= 1.0, (
                f"{algo} cv_score {tm.cv_score} out of [0, 1]"
            )

    def test_metrics_accuracy_in_unit_interval(self, trainer, dataset_csv):
        results = trainer.train_all(dataset_csv)
        for algo, tm in results.items():
            assert 0.0 <= tm.metrics.accuracy <= 1.0, (
                f"{algo} accuracy {tm.metrics.accuracy} out of [0, 1]"
            )

    def test_preprocessor_artifact_created(self, trainer, dataset_csv, tmp_dirs):
        trainer.train_all(dataset_csv)
        assert os.path.exists(tmp_dirs["preprocessor_path"]), (
            "Preprocessor pipeline artifact was not created."
        )

    def test_feature_mask_artifact_created(self, trainer, dataset_csv, tmp_dirs):
        trainer.train_all(dataset_csv)
        assert os.path.exists(tmp_dirs["mask_path"]), (
            "Feature mask artifact was not created."
        )


# ---------------------------------------------------------------------------
# Tests: train_all with DB repo (mocked)
# ---------------------------------------------------------------------------

class TestTrainAllWithRepo:
    """Verify DB interactions when a model_log_repo is provided."""

    def _make_mock_repo(self):
        """Return a mock ModelLogRepository whose save() returns its argument."""
        repo = MagicMock()
        # save() should return the log object it receives (simulates DB assign of id)
        saved_counter = [0]

        def _save(log):
            saved_counter[0] += 1
            log.id = saved_counter[0]
            return log

        repo.save.side_effect = _save
        return repo

    def test_save_called_at_least_four_times(self, trainer, dataset_csv):
        """save() is called once per algorithm (initial save) plus updates."""
        repo = self._make_mock_repo()
        trainer.train_all(dataset_csv, model_log_repo=repo)
        # At minimum 4 initial saves (one per algorithm)
        assert repo.save.call_count >= 4

    def test_save_called_for_each_algorithm(self, trainer, dataset_csv):
        """Each of the four algorithms produces at least one save() call."""
        repo = self._make_mock_repo()
        trainer.train_all(dataset_csv, model_log_repo=repo)

        # Collect algorithm names from all save() calls
        saved_algorithms = [
            call_args.args[0].algorithm
            for call_args in repo.save.call_args_list
        ]
        for algo in ["logistic_regression", "decision_tree", "random_forest", "svm"]:
            assert algo in saved_algorithms, (
                f"No save() call found for algorithm '{algo}'"
            )

    def test_best_model_marked_deployed(self, trainer, dataset_csv):
        """The algorithm with the highest test accuracy has deployed=True."""
        repo = self._make_mock_repo()
        results = trainer.train_all(dataset_csv, model_log_repo=repo)

        # Find the best algorithm by accuracy
        best_algo = max(results, key=lambda a: results[a].metrics.accuracy)

        # Find the final state of each algorithm's log from save() calls
        # The last save() call for the best algorithm should have deployed=True
        final_states: dict[str, bool] = {}
        for call_args in repo.save.call_args_list:
            log = call_args.args[0]
            final_states[log.algorithm] = log.deployed

        assert final_states.get(best_algo) is True, (
            f"Best algorithm '{best_algo}' was not marked as deployed. "
            f"Final states: {final_states}"
        )

    def test_non_best_models_not_deployed(self, trainer, dataset_csv):
        """Only the best model should be marked deployed=True."""
        repo = self._make_mock_repo()
        results = trainer.train_all(dataset_csv, model_log_repo=repo)

        best_algo = max(results, key=lambda a: results[a].metrics.accuracy)

        # Collect final deployed state per algorithm
        final_states: dict[str, bool] = {}
        for call_args in repo.save.call_args_list:
            log = call_args.args[0]
            final_states[log.algorithm] = log.deployed

        for algo, deployed in final_states.items():
            if algo != best_algo:
                assert deployed is False, (
                    f"Algorithm '{algo}' should not be deployed but has deployed={deployed}"
                )

    def test_feature_set_stored_as_json(self, trainer, dataset_csv):
        """feature_set column in saved logs must be valid JSON."""
        repo = self._make_mock_repo()
        trainer.train_all(dataset_csv, model_log_repo=repo)

        for call_args in repo.save.call_args_list:
            log = call_args.args[0]
            # Should not raise
            parsed = json.loads(log.feature_set)
            assert isinstance(parsed, list), (
                f"feature_set for {log.algorithm} is not a JSON list"
            )
            assert len(parsed) > 0, (
                f"feature_set for {log.algorithm} is empty"
            )

    def test_metrics_stored_correctly(self, trainer, dataset_csv):
        """Accuracy stored in the log matches the TrainedModel metrics."""
        repo = self._make_mock_repo()
        results = trainer.train_all(dataset_csv, model_log_repo=repo)

        # Build a map of algorithm → first save() call (initial metrics save)
        first_saves: dict[str, object] = {}
        for call_args in repo.save.call_args_list:
            log = call_args.args[0]
            if log.algorithm not in first_saves:
                first_saves[log.algorithm] = log

        for algo, log in first_saves.items():
            expected_accuracy = results[algo].metrics.accuracy
            assert log.accuracy == pytest.approx(expected_accuracy, abs=1e-6), (
                f"{algo}: stored accuracy {log.accuracy} != "
                f"computed accuracy {expected_accuracy}"
            )


# ---------------------------------------------------------------------------
# Tests: error handling
# ---------------------------------------------------------------------------

class TestTrainAllErrors:
    def test_raises_file_not_found_for_missing_csv(self, trainer):
        with pytest.raises(FileNotFoundError, match="Dataset CSV not found"):
            trainer.train_all("/nonexistent/path/dataset.csv")

    def test_raises_value_error_for_missing_columns(self, trainer, tmp_dirs):
        """A CSV missing required columns raises ValueError."""
        bad_csv = os.path.join(tmp_dirs["artifacts_dir"], "bad.csv")
        pd.DataFrame({"col_a": [1, 2], "col_b": [3, 4]}).to_csv(bad_csv, index=False)
        with pytest.raises(ValueError, match="missing required columns"):
            trainer.train_all(bad_csv)