tests/test_data.py

import pytest
import pandas as pd
import numpy as np
from src.data import generate_synthetic_data, Preprocessor, prepare_data

def test_generate_synthetic_data():
    n_samples = 150
    df = generate_synthetic_data(n_samples=n_samples, random_state=42)
    
    assert isinstance(df, pd.DataFrame)
    assert len(df) == n_samples
    assert list(df.columns) == ['age', 'monthly_charges', 'contract_length', 'support_calls', 'tech_support', 'churn']
    assert df['churn'].isin([0, 1]).all()
    assert df['tech_support'].isin(['yes', 'no']).all()

def test_preprocessor():
    df = pd.DataFrame({
        'age': [20, 40, 60],
        'monthly_charges': [30.0, 75.0, 110.0],
        'contract_length': [1, 12, 24],
        'support_calls': [0, 3, 5],
        'tech_support': ['yes', 'no', 'yes'],
        'churn': [0, 1, 0]
    })
    
    preprocessor = Preprocessor()
    X_trans = preprocessor.fit_transform(df)
    
    # 4 numerical variables + 1 binary encoded categorical = 5 columns
    assert X_trans.shape == (3, 5)
    
    # Verify that fit sets the is_fitted flag
    assert preprocessor.is_fitted is True
    
    # Check that calling transform without fitting raises ValueError on a new instance
    unfitted = Preprocessor()
    with pytest.raises(ValueError):
        unfitted.transform(df)

def test_prepare_data():
    df = generate_synthetic_data(n_samples=200, random_state=42)
    X_train, X_test, y_train, y_test, preprocessor = prepare_data(df, test_size=0.2, random_state=42)
    
    assert X_train.shape[0] == 160
    assert X_test.shape[0] == 40
    assert X_train.shape[1] == 5
    assert len(y_train) == 160
    assert len(y_test) == 40
    assert isinstance(preprocessor, Preprocessor)