DataPrep/tests/test_transformation.py at main · Medalcode/DataPrep · GitHub

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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
"""
tests/test_transformation.py
-----------------------------
Unit tests for the transformation module.
"""

import sys
from pathlib import Path

import pandas as pd
import pytest

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from src.transformation import (
    add_derived_columns,
    normalize_numeric_columns,
    parse_date_columns,
    transform_data,
)


@pytest.fixture
def base_df():
    return pd.DataFrame(
        {
            "fecha": ["2023-01-15", "2023-06-20", "2024-03-10"],
            "producto": ["Laptop", "Mouse", "Teclado"],
            "precio": [999.99, 29.99, 59.99],
            "cantidad": [1, 5, 3],
        }
    )


class TestParseDateColumns:
    def test_converts_to_datetime(self, base_df):
        result = parse_date_columns(base_df)
        assert pd.api.types.is_datetime64_any_dtype(result["fecha"])

    def test_extracts_year(self, base_df):
        result = parse_date_columns(base_df)
        assert "fecha_year" in result.columns
        assert result["fecha_year"].iloc[0] == 2023

    def test_extracts_month(self, base_df):
        result = parse_date_columns(base_df)
        assert "fecha_month" in result.columns
        assert result["fecha_month"].iloc[0] == 1

    def test_extracts_day(self, base_df):
        result = parse_date_columns(base_df)
        assert "fecha_day" in result.columns
        assert result["fecha_day"].iloc[0] == 15

    def test_handles_bad_date_gracefully(self):
        df = pd.DataFrame({"fecha": ["2023-01-15", "not-a-date", "2024-12-01"]})
        result = parse_date_columns(df, date_cols=["fecha"])
        assert pd.api.types.is_datetime64_any_dtype(result["fecha"])
        assert result["fecha"].isna().sum() == 1  # bad date becomes NaT


class TestAddDerivedColumns:
    def test_creates_total_column(self, base_df):
        result = add_derived_columns(base_df)
        assert "total" in result.columns

    def test_total_is_precio_times_cantidad(self, base_df):
        result = add_derived_columns(base_df)
        assert result["total"].iloc[0] == pytest.approx(999.99 * 1, abs=0.01)

    def test_creates_iva_column(self, base_df):
        result = add_derived_columns(base_df)
        assert "total_con_iva" in result.columns

    def test_total_con_iva_correct(self, base_df):
        result = add_derived_columns(base_df)
        expected = result["total"].iloc[0] * 1.19
        assert result["total_con_iva"].iloc[0] == pytest.approx(expected, abs=0.01)

    def test_creates_category_column(self, base_df):
        result = add_derived_columns(base_df)
        assert "categoria_venta" in result.columns


class TestNormalizeNumericColumns:
    def test_creates_norm_column(self, base_df):
        result = normalize_numeric_columns(base_df, columns=["precio"])
        assert "precio_norm" in result.columns

    def test_norm_values_between_0_and_1(self, base_df):
        result = normalize_numeric_columns(base_df, columns=["precio"])
        assert result["precio_norm"].between(0, 1).all()

    def test_zscore_normalization(self, base_df):
        result = normalize_numeric_columns(base_df, columns=["precio"], method="zscore")
        assert "precio_norm" in result.columns


class TestTransformData:
    def test_returns_dataframe(self, base_df):
        result = transform_data(base_df)
        assert isinstance(result, pd.DataFrame)

    def test_has_more_columns_than_input(self, base_df):
        result = transform_data(base_df)
        assert len(result.columns) > len(base_df.columns)

    def test_runs_without_error(self, base_df):
        # Should not raise
        transform_data(base_df, normalize=True)