iBrokeTheCode/Multimodal_Product_Classification

1from unittest.mock import patch2 3import pytest4from sklearn.datasets import make_classification5from sklearn.decomposition import PCA6from sklearn.ensemble import RandomForestClassifier7from sklearn.linear_model import LogisticRegression8from sklearn.model_selection import train_test_split9 10from src.classifiers_classic_ml import train_and_evaluate_model, visualize_embeddings11 12####################################################################################################13################################### Test the Classical ML Models ###################################14####################################################################################################15 16 17@pytest.fixture18def sample_embedding_data():19    """20    Fixture to create a mock dataset for testing dimensionality reduction and model training.21    Returns:22        X_train, X_test, y_train, y_test: Training and testing data along with labels.23    """24    # Create a synthetic dataset with 20 samples, 6 features, and 3 classes25    X, y = make_classification(26        n_samples=20, n_features=6, n_classes=3, random_state=42, n_informative=427    )28 29    # Split the dataset into training and test sets (80% train, 20% test)30    X_train, X_test, y_train, y_test = train_test_split(31        X, y, test_size=0.2, random_state=4232    )33 34    return X_train, X_test, y_train, y_test35 36 37@pytest.mark.parametrize(38    "method, plot_type",39    [40        ("PCA", "2D"),  # PCA reduction to 2D41        ("PCA", "3D"),  # PCA reduction to 3D42    ],43)44def test_visualize_embeddings(method, plot_type, sample_embedding_data):45    """46    Test the dimensionality reduction and embedding visualization.47    This ensures that PCA can reduce embeddings correctly and produce visualizations.48    """49    X_train, X_test, y_train, y_test = sample_embedding_data50 51    # Mock the plotly figures to avoid actual plotting in test environment52    with patch("plotly.graph_objs.Figure.show"):53        # Test the visualize_embeddings function54        model = visualize_embeddings(55            X_train, X_test, y_train, y_test, plot_type=plot_type, method=method56        )57 58    # Check if the PCA model is an instance of the correct class and has the expected number of components59    assert isinstance(model, PCA), "The model should be an instance of PCA"60    if plot_type == "2D":61        assert model.n_components_ == 2, "PCA should reduce data to 2 components"62    elif plot_type == "3D":63        assert model.n_components_ == 3, "PCA should reduce data to 3 components"64 65 66def test_train_and_evaluate_model(sample_embedding_data):67    """68    Test the training and evaluation of models (Logistic Regression, Random Forest).69    Ensures that models are correctly trained and returned in the expected format.70    """71    X_train, X_test, y_train, y_test = sample_embedding_data72 73    # Train and evaluate the models74    trained_models = train_and_evaluate_model(75        X_train, X_test, y_train, y_test, test=False76    )77 78    # Verify that trained_models is a list79    assert isinstance(trained_models, list), (80        "The output should be a list of trained models"81    )82 83    # Check that at least two models were trained (Logistic Regression, Random Forest)84    assert len(trained_models) >= 2, "At least two models should be trained"85 86    # Check that the models have Logistic Regression and Random Forest87    models_instances = [model for _, model in trained_models]88    assert any(isinstance(model, LogisticRegression) for model in models_instances), (89        "Logistic Regression model not found"90    )91    assert any(92        isinstance(model, RandomForestClassifier) for model in models_instances93    ), "Random Forest model not found"94 95    # Ensure that the trained models are indeed fitted (trained)96    for name, model in trained_models:97        assert hasattr(model, "fit"), f"{name} should have a fit method"98        assert hasattr(model, "predict"), f"{name} should have a predict method"99 100        # Check if the model is correctly trained by predicting on the test set101        y_pred = model.predict(X_test)102        assert y_pred is not None, f"{name} should have successfully made predictions"103 104 105if __name__ == "__main__":106    pytest.main()107