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!!