File size: 1,717 Bytes
d0bac84 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | 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)
|