"""Tests for mldash._serialize — artifact serialization."""
import os

import joblib
import pytest

from mldash._serialize import serialize_artifacts


class TestSerializeArtifacts:
    def test_serializes_model_only(self, tmp_dir, fitted_classifier):
        paths = serialize_artifacts(tmp_dir, fitted_classifier)
        assert 'model.pkl' in paths
        assert os.path.isfile(paths['model.pkl'])
        # Should be loadable
        loaded = joblib.load(paths['model.pkl'])
        assert type(loaded).__name__ == 'LogisticRegression'

    def test_does_not_include_optional_when_none(self, tmp_dir, fitted_classifier):
        paths = serialize_artifacts(tmp_dir, fitted_classifier)
        assert 'scaler.pkl' not in paths
        assert 'encoder.pkl' not in paths
        assert 'label_map.pkl' not in paths

    def test_serializes_with_scaler(self, tmp_dir, fitted_classifier, fitted_scaler):
        paths = serialize_artifacts(tmp_dir, fitted_classifier, scaler=fitted_scaler)
        assert 'model.pkl' in paths
        assert 'scaler.pkl' in paths
        assert os.path.isfile(paths['scaler.pkl'])
        loaded = joblib.load(paths['scaler.pkl'])
        assert hasattr(loaded, 'transform')

    def test_serializes_with_encoder(self, tmp_dir, fitted_classifier, fitted_encoder):
        paths = serialize_artifacts(tmp_dir, fitted_classifier, encoder=fitted_encoder)
        assert 'encoder.pkl' in paths
        assert os.path.isfile(paths['encoder.pkl'])

    def test_serializes_with_label_map(self, tmp_dir, fitted_classifier, sample_label_map):
        paths = serialize_artifacts(tmp_dir, fitted_classifier, label_map=sample_label_map)
        assert 'label_map.pkl' in paths
        loaded = joblib.load(paths['label_map.pkl'])
        assert loaded == sample_label_map

    def test_serializes_all_artifacts(self, tmp_dir, fitted_classifier,
                                      fitted_scaler, fitted_encoder, sample_label_map):
        paths = serialize_artifacts(
            tmp_dir, fitted_classifier,
            scaler=fitted_scaler,
            encoder=fitted_encoder,
            label_map=sample_label_map,
        )
        assert len(paths) == 4
        for name in ('model.pkl', 'scaler.pkl', 'encoder.pkl', 'label_map.pkl'):
            assert name in paths
            assert os.path.isfile(paths[name])
