Best Python code snippet using pandera_python
test_model.py
Source:test_model.py  
...176    """Test validate method on valid data."""177    class Schema(pa.SchemaModel):178        a: Series[int]179        @pa.check("a")180        def int_column_lt_100(cls, series: pd.Series) -> Iterable[bool]:181            # pylint:disable=no-self-argument182            assert cls is Schema183            return series < 100184    df = pd.DataFrame({"a": [99]})185    assert isinstance(Schema.validate(df, lazy=True), pd.DataFrame)186def test_check_validate_method_field() -> None:187    """Test validate method on valid data."""188    class Schema(pa.SchemaModel):189        a: Series[int] = pa.Field()190        b: Series[int]191        @pa.check(a)192        def int_column_lt_200(cls, series: pd.Series) -> Iterable[bool]:193            # pylint:disable=no-self-argument194            assert cls is Schema195            return series < 200196        @pa.check(a, "b")197        def int_column_lt_100(cls, series: pd.Series) -> Iterable[bool]:198            # pylint:disable=no-self-argument199            assert cls is Schema200            return series < 100201    df = pd.DataFrame({"a": [99], "b": [99]})202    assert isinstance(Schema.validate(df, lazy=True), pd.DataFrame)203def test_check_validate_method_aliased_field() -> None:204    """Test validate method on valid data."""205    class Schema(pa.SchemaModel):206        a: Series[int] = pa.Field(alias=2020, gt=50)207        @pa.check(a)208        def int_column_lt_100(cls, series: pd.Series) -> Iterable[bool]:209            # pylint:disable=no-self-argument210            assert cls is Schema211            return series < 100212    df = pd.DataFrame({2020: [99]})213    assert len(Schema.to_schema().columns[2020].checks) == 2214    assert isinstance(Schema.validate(df, lazy=True), pd.DataFrame)215def test_check_single_column() -> None:216    """Test the behaviour of a check on a single column."""217    class Schema(pa.SchemaModel):218        a: Series[int]219        @pa.check("a")220        def int_column_lt_100(cls, series: pd.Series) -> Iterable[bool]:221            # pylint:disable=no-self-argument222            assert cls is Schema223            return series < 100224    df = pd.DataFrame({"a": [101]})225    schema = Schema.to_schema()226    err_msg = r"Column\s*a\s*int_column_lt_100\s*\[101\]\s*1"227    with pytest.raises(pa.errors.SchemaErrors, match=err_msg):228        schema.validate(df, lazy=True)229def test_check_single_index() -> None:230    """Test the behaviour of a check on a single index."""231    class Schema(pa.SchemaModel):232        a: Index[str]233        @pa.check("a")234        def not_dog(cls, idx: pd.Index) -> Iterable[bool]:235            # pylint:disable=no-self-argument236            assert cls is Schema237            return ~idx.str.contains("dog")238    df = pd.DataFrame(index=["cat", "dog"])239    err_msg = r"Index\s*<NA>\s*not_dog\s*\[dog\]\s*"240    with pytest.raises(pa.errors.SchemaErrors, match=err_msg):241        Schema.validate(df, lazy=True)242def test_field_and_check() -> None:243    """Test the combination of a field and a check on the same column."""244    class Schema(pa.SchemaModel):245        a: Series[int] = pa.Field(eq=1)246        @pa.check("a")247        @classmethod248        def int_column_lt_100(cls, series: pd.Series) -> Iterable[bool]:249            return series < 100250    schema = Schema.to_schema()251    assert len(schema.columns["a"].checks) == 2252def test_check_non_existing() -> None:253    """Test a check on a non-existing column."""254    class Schema(pa.SchemaModel):255        a: Series[int]256        @pa.check("nope")257        @classmethod258        def int_column_lt_100(cls, series: pd.Series) -> Iterable[bool]:259            return series < 100260    err_msg = (261        "Check int_column_lt_100 is assigned to a non-existing field 'nope'"262    )263    with pytest.raises(pa.errors.SchemaInitError, match=err_msg):264        Schema.to_schema()265def test_multiple_checks() -> None:266    """Test multiple checks on the same column."""267    class Schema(pa.SchemaModel):268        a: Series[int]269        @pa.check("a")270        @classmethod271        def int_column_lt_100(cls, series: pd.Series) -> Iterable[bool]:272            return series < 100273        @pa.check("a")274        @classmethod275        def int_column_gt_0(cls, series: pd.Series) -> Iterable[bool]:276            return series > 0277    schema = Schema.to_schema()278    assert len(schema.columns["a"].checks) == 2279    df = pd.DataFrame({"a": [0]})280    err_msg = r"Column\s*a\s*int_column_gt_0\s*\[0\]\s*1"281    with pytest.raises(pa.errors.SchemaErrors, match=err_msg):282        schema.validate(df, lazy=True)283    df = pd.DataFrame({"a": [101]})284    err_msg = r"Column\s*a\s*int_column_lt_100\s*\[101\]\s*1"285    with pytest.raises(pa.errors.SchemaErrors, match=err_msg):286        schema.validate(df, lazy=True)287def test_check_multiple_columns() -> None:288    """Test a single check decorator targeting multiple columns."""289    class Schema(pa.SchemaModel):290        a: Series[int]291        b: Series[int]292        @pa.check("a", "b")293        @classmethod294        def int_column_lt_100(cls, series: pd.Series) -> Iterable[bool]:295            return series < 100296    df = pd.DataFrame({"a": [101], "b": [200]})297    with pytest.raises(298        pa.errors.SchemaErrors, match="2 schema errors were found"299    ):300        Schema.validate(df, lazy=True)301def test_check_regex() -> None:302    """Test the regex argument of the check decorator."""303    class Schema(pa.SchemaModel):304        a: Series[int]305        abc: Series[int]306        cba: Series[int]307        @pa.check("^a", regex=True)308        @classmethod309        def int_column_lt_100(cls, series: pd.Series) -> Iterable[bool]:310            return series < 100311    df = pd.DataFrame({"a": [101], "abc": [1], "cba": [200]})312    with pytest.raises(313        pa.errors.SchemaErrors, match="1 schema errors were found"314    ):315        Schema.validate(df, lazy=True)316def test_inherit_schemamodel_fields() -> None:317    """Test that columns and indices are inherited."""318    class Base(pa.SchemaModel):319        a: Series[int]320        idx: Index[str]321    class Mid(Base):322        b: Series[str]323        idx: Index[str]324    class Child(Mid):325        b: Series[int]326    expected = pa.DataFrameSchema(327        columns={"a": pa.Column(int), "b": pa.Column(int)},328        index=pa.Index(str),329    )330    assert expected == Child.to_schema()331def test_inherit_schemamodel_fields_alias() -> None:332    """Test that columns and index aliases are inherited."""333    class Base(pa.SchemaModel):334        a: Series[int]335        idx: Index[str]336    class Mid(Base):337        b: Series[str] = pa.Field(alias="_b")338        idx: Index[str]339    class ChildOverrideAttr(Mid):340        b: Series[int]341    class ChildOverrideAlias(Mid):342        b: Series[str] = pa.Field(alias="new_b")343    class ChildNewAttr(Mid):344        c: Series[int]345    class ChildEmpty(Mid):346        pass347    expected_mid = pa.DataFrameSchema(348        columns={"a": pa.Column(int), "_b": pa.Column(str)},349        index=pa.Index(str),350    )351    expected_child_override_attr = expected_mid.rename_columns(352        {"_b": "b"}353    ).update_column("b", dtype=int)354    expected_child_override_alias = expected_mid.rename_columns(355        {"_b": "new_b"}356    )357    expected_child_new_attr = expected_mid.add_columns(358        {359            "c": pa.Column(int),360        }361    )362    assert expected_mid == Mid.to_schema()363    assert expected_child_override_attr == ChildOverrideAttr.to_schema()364    assert expected_child_override_alias == ChildOverrideAlias.to_schema()365    assert expected_child_new_attr == ChildNewAttr.to_schema()366    assert expected_mid == ChildEmpty.to_schema()367def test_inherit_field_checks() -> None:368    """Test that checks are inherited and overridden."""369    class Base(pa.SchemaModel):370        a: Series[int]371        abc: Series[int]372        @pa.check("^a", regex=True)373        @classmethod374        def a_max(cls, series: pd.Series) -> Iterable[bool]:375            return series < 100376        @pa.check("a")377        @classmethod378        def a_min(cls, series: pd.Series) -> Iterable[bool]:379            return series > 1380    class Child(Base):381        @pa.check("a")382        @classmethod383        def a_max(cls, series: pd.Series) -> Iterable[bool]:384            return series < 10385    schema = Child.to_schema()386    assert len(schema.columns["a"].checks) == 2387    assert len(schema.columns["abc"].checks) == 0388    df = pd.DataFrame({"a": [15], "abc": [100]})389    err_msg = r"Column\s*a\s*a_max\s*\[15\]\s*1"390    with pytest.raises(pa.errors.SchemaErrors, match=err_msg):391        schema.validate(df, lazy=True)392def test_dataframe_check() -> None:393    """Test dataframe checks."""394    class Base(pa.SchemaModel):395        a: Series[int]396        b: Series[int]397        @pa.dataframe_check398        @classmethod399        def value_max(cls, df: pd.DataFrame) -> Iterable[bool]:400            return df < 200401    class Child(Base):402        @pa.dataframe_check()403        @classmethod404        def value_min(cls, df: pd.DataFrame) -> Iterable[bool]:405            return df > 0406        @pa.dataframe_check407        @classmethod408        def value_max(cls, df: pd.DataFrame) -> Iterable[bool]:409            return df < 100410    schema = Child.to_schema()411    assert len(schema.checks) == 2412    df = pd.DataFrame({"a": [101, 1], "b": [1, 0]})413    with pytest.raises(414        pa.errors.SchemaErrors, match="2 schema errors were found"415    ):416        schema.validate(df, lazy=True)417def test_registered_dataframe_checks(418    extra_registered_checks: None,  # pylint: disable=unused-argument419) -> None:420    """Check that custom check inheritance works"""421    # pylint: disable=unused-variable422    @pax.register_check_method(statistics=["one_arg"])423    def base_check(df, *, one_arg):424        # pylint: disable=unused-argument425        return True426    @pax.register_check_method(statistics=["one_arg", "two_arg"])427    def child_check(df, *, one_arg, two_arg):428        # pylint: disable=unused-argument429        return True430    # pylint: enable=unused-variable431    check_vals = {432        "one_arg": 150,433        "two_arg": "hello",434        "one_arg_prime": "not_150",435    }436    class Base(pa.SchemaModel):437        a: Series[int]438        b: Series[int]439        class Config:440            no_param_check = ()441            base_check = check_vals["one_arg"]442    class Child(Base):443        class Config:444            base_check = check_vals["one_arg_prime"]445            child_check = {446                "one_arg": check_vals["one_arg"],447                "two_arg": check_vals["two_arg"],448            }449    base = Base.to_schema()450    child = Child.to_schema()451    expected_stats_base = {452        "no_param_check": {},453        "base_check": {"one_arg": check_vals["one_arg"]},454    }455    expected_stats_child = {456        "no_param_check": {},457        "base_check": {"one_arg": check_vals["one_arg_prime"]},458        "child_check": {459            "one_arg": check_vals["one_arg"],460            "two_arg": check_vals["two_arg"],461        },462    }463    assert {b.name: b.statistics for b in base.checks} == expected_stats_base464    assert {c.name: c.statistics for c in child.checks} == expected_stats_child465    # check that unregistered checks raise466    with pytest.raises(AttributeError, match=".*custom checks.*"):467        class ErrorSchema(pa.SchemaModel):468            class Config:469                unknown_check = {}  # type: ignore[var-annotated]470        # Check lookup happens at validation/to_schema conversion time471        # This means that you can register checks after defining a Config,472        # but also because of caching you can refer to a check that no longer473        # exists for some order of operations.474        ErrorSchema.to_schema()475def test_config() -> None:476    """Test that Config can be inherited and translate into DataFrameSchema options."""477    class Base(pa.SchemaModel):478        a: Series[int]479        idx_1: Index[str]480        idx_2: Index[str]481        class Config:482            name = "Base schema"483            coerce = True484            ordered = True485            multiindex_coerce = True486            multiindex_strict = True487            multiindex_name: Optional[str] = "mi"488    class Child(Base):489        b: Series[int]490        class Config:491            name = "Child schema"492            strict = True493            multiindex_strict = False494            description = "foo"495            title = "bar"496    expected = pa.DataFrameSchema(497        columns={"a": pa.Column(int), "b": pa.Column(int)},498        index=pa.MultiIndex(499            [pa.Index(str, name="idx_1"), pa.Index(str, name="idx_2")],500            coerce=True,501            strict=False,502            name="mi",503        ),504        name="Child schema",505        coerce=True,506        strict=True,507        ordered=True,508        description="foo",509        title="bar",510    )511    assert expected == Child.to_schema()512def test_config_docstrings() -> None:513    class Model(pa.SchemaModel):514        """foo"""515        a: Series[int]516    assert Model.__doc__ == Model.to_schema().description517class Input(pa.SchemaModel):518    a: Series[int]519    b: Series[int]520    idx: Index[str]521class Output(Input):522    c: Series[int]523def test_check_types() -> None:524    @pa.check_types525    def transform(df: DataFrame[Input]) -> DataFrame[Output]:526        return df.assign(c=lambda x: x.a + x.b)527    data = pd.DataFrame(528        {"a": [1, 2, 3], "b": [4, 5, 6]}, index=pd.Index(["a", "b", "c"])529    )530    assert isinstance(transform(data), pd.DataFrame)531    for invalid_data in [532        data.drop("a", axis="columns"),533        data.drop("b", axis="columns"),534        data.assign(a=["a", "b", "c"]),535        data.assign(b=["a", "b", "c"]),536        data.reset_index(drop=True),537    ]:538        with pytest.raises(pa.errors.SchemaError):539            transform(invalid_data)540def test_alias() -> None:541    """Test that columns and indices can be aliased."""542    class Schema(pa.SchemaModel):543        col_2020: Series[int] = pa.Field(alias=2020)544        idx: Index[int] = pa.Field(alias="_idx", check_name=True)545        @pa.check(2020)546        @classmethod547        def int_column_lt_100(cls, series: pd.Series) -> Iterable[bool]:548            return series < 100549    schema = Schema.to_schema()550    assert len(schema.columns) == 1551    assert schema.columns.get(2020, None) is not None552    assert schema.index.name == "_idx"553    df = pd.DataFrame({2020: [99]}, index=[0])554    df.index.name = "_idx"555    assert len(Schema.to_schema().columns[2020].checks) == 1556    assert isinstance(Schema.validate(df), pd.DataFrame)557    # test multiindex558    class MISchema(pa.SchemaModel):559        idx1: Index[int] = pa.Field(alias="index0")560        idx2: Index[int] = pa.Field(alias="index1")561    actual = [index.name for index in MISchema.to_schema().index.indexes]...Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.
You could also refer to video tutorials over LambdaTest YouTube channel to get step by step demonstration from industry experts.
Get 100 minutes of automation test minutes FREE!!
