tests/test_utils.py

"""
Testes simples para as funções de utilidade.
"""

import pandas as pd

from src.utils import preparar_base


def test_preparar_base_basico():
    """
    Testa se a função preparar_base retorna shapes válidos.
    """
    # Pequeno DataFrame de exemplo
    dados = {
        "ID": [1, 2, 3, 4],
        "Income": [1000.0, 2000.0, 1500.0, 1800.0],
        "Complain": [0, 1, 0, 1],
        "Dt_Customer": ["01-01-2020", "05-01-2020", "10-01-2020", "15-01-2020"],
    }
    df = pd.DataFrame(dados)

    X_train, X_test, y_train, y_test, features, scaler = preparar_base(
        df, aplicar_smote=False, teste=0.5, random_state=42
    )

    assert X_train.shape[0] == y_train.shape[0]
    assert X_test.shape[0] == y_test.shape[0]
    assert len(features) > 0