activitysim/activitysim/abm/models/util/test/test_input_checker.py at 15f77d7e48e34382f4ffa2f6acb50c37dbc93122 · ActivitySim/activitysim · 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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
# ActivitySim
# See full license in LICENSE.txt.

from __future__ import annotations

import os.path

import pandas as pd
import pandas.testing as pdt
import pandera.pandas as pa
import pytest
import yaml

from activitysim.abm.models.input_checker import TABLE_STORE, validate_with_pandera


@pytest.fixture(scope="class")
def v_errors():
    v_errors = {}
    v_errors["households"] = []
    return v_errors


@pytest.fixture(scope="class")
def v_warnings():
    v_warnings = {}
    v_warnings["households"] = []
    return v_warnings


@pytest.fixture(scope="module")
def validation_settings():
    return {"method": "pandera", "class": "Household"}


@pytest.fixture(scope="module")
def households():
    return pd.DataFrame(
        data={
            "household_id": [1, 2, 3, 4],
            "home_zone_id": [0, 1, 2, 3],
            "income": [10000, 20000, 30000, 40000],
            "hhsize": [2, 2, 4, 5],
        }
    )


TABLE_STORE["households"] = pd.DataFrame(
    data={
        "household_id": [1, 2, 3, 4],
        "home_zone_id": [0, 1, 2, 3],
        "income": [10000, 20000, 30000, 40000],
        "hhsize": [2, 2, 4, 5],
    }
)


def test_passing_dataframe(households, v_errors, v_warnings, validation_settings):
    TABLE_STORE["households"] = households

    class input_checker:
        class Household(pa.DataFrameModel):
            household_id: int = pa.Field(unique=True, gt=0)
            home_zone_id: int = pa.Field(ge=0)
            hhsize: int = pa.Field(gt=0)
            income: int = pa.Field(ge=0, raise_warning=True)

        @pa.dataframe_check(name="Example setup of a passing error check.")
        def dummy_example(cls, households: pd.DataFrame):
            return (households.household_id > 0).all()

    returned_errors, returned_warnings = validate_with_pandera(
        input_checker, "households", validation_settings, v_errors, v_warnings
    )

    assert (
        len(returned_errors["households"]) == 0
    ), f"Expect no household errors, but got {returned_errors}"
    assert (
        len(returned_warnings["households"]) == 0
    ), f"Expect no household warnings, but got {returned_warnings}"


def test_error_dataframe(households, v_errors, v_warnings, validation_settings):
    TABLE_STORE["households"] = households

    class input_checker:
        class Household(pa.DataFrameModel):
            household_id: int = pa.Field(unique=True, gt=0)
            home_zone_id: int = pa.Field(ge=0)
            hhsize: int = pa.Field(gt=0)
            income: int = pa.Field(ge=0, raise_warning=True)
            bug1: int  # error here

    returned_errors, returned_warnings = validate_with_pandera(
        input_checker, "households", validation_settings, v_errors, v_warnings
    )

    assert (
        len(returned_errors["households"]) == 1
    ), f"Expected household error, but got {len(returned_errors['households'])}"
    assert (
        len(returned_warnings["households"]) == 0
    ), f"Expect no household warnings, but got {returned_warnings['households']}"


def test_warning_dataframe(households, v_errors, v_warnings, validation_settings):
    TABLE_STORE["households"] = households

    class input_checker:
        class Household(pa.DataFrameModel):
            household_id: int = pa.Field(unique=True, gt=0)
            home_zone_id: int = pa.Field(ge=0)
            hhsize: int = pa.Field(gt=0)
            income: int = pa.Field(ge=100000, raise_warning=True)  # warning here

    returned_errors, returned_warnings = validate_with_pandera(
        input_checker, "households", validation_settings, v_errors, v_warnings
    )

    assert (
        len(returned_errors["households"]) == 0
    ), f"Expect no household errors, but got {returned_errors['households']}"
    assert (
        len(returned_warnings["households"]) > 0
    ), f"Expected warnings, but got {len(returned_warnings['households'])}"


def test_custom_check_failure_dataframe(
    households, v_errors, v_warnings, validation_settings
):
    TABLE_STORE["households"] = households

    class input_checker:
        class Household(pa.DataFrameModel):
            household_id: int = pa.Field(unique=True, gt=0)
            home_zone_id: int = pa.Field(ge=0)
            hhsize: int = pa.Field(gt=0)
            income: int = pa.Field(ge=0, raise_warning=True)

            @pa.dataframe_check(name="Example setup of a failed error check.")
            def dummy_example(cls, households: pd.DataFrame):
                return False

    returned_errors, returned_warnings = validate_with_pandera(
        input_checker, "households", validation_settings, v_errors, v_warnings
    )

    assert (
        len(returned_errors["households"]) > 0
    ), f"Expected household errors, but got {returned_errors['households']}"
    assert (
        len(returned_warnings["households"]) == 0
    ), f"Expect no household warnings, but got {returned_warnings['households']}"