edit functools example

Author: khuyentran1401

data_science_tools/pandera_hypothesis/test1.py

@@ -10,10 +10,13 @@ def processing_fn(df):
 def test_processing_fn():
     # Create test data
     df = pd.DataFrame({"val1": [1, 1, -1, -2, 2], "val2": [1, 2, -2, -1, 2]})
+
     # Get result
     result = processing_fn(df)
+
     # Create expected output
     expected = df.copy()
     expected["val3"] = [1, 0.5, 0.5, 2, 1]
+    
     # Test
     assert_frame_equal(result, expected, check_dtype=False)

data_science_tools/pandera_hypothesis/test3.py

@@ -1,7 +1,9 @@
 import pandas as pd
+import pytest
+from pandas.testing import assert_frame_equal
 import pandera as pa
 
-out_schema = pa.DataFrameSchema(
+expected = pa.DataFrameSchema(
     {
         "val1": pa.Column(int, pa.Check.in_range(-2, 3)),
         "val2": pa.Column(int, pa.Check.in_range(-2, 3)),
@@ -10,12 +12,20 @@
 )
 
 
-@pa.check_output(out_schema)
+@pa.check_output(expected)
 def processing_fn(df):
     processed = df.assign(val3=df.val1 / df.val2)
     return processed
 
 
-if __name__ == "__main__":
-    df = pd.DataFrame({"val1": [1, 1, -1, -2, 2], "val2": [1, 1, -1, -2, 2]})
-    processing_fn(df)
+val1 = [[1, 1, -1, -2, 2], [1, 1, -1, -2, 2]]
+val2 = [[1, 2, -2, -1, 2], [1, 1, 1, 1, 1]]
+
+
+@pytest.mark.parametrize("val1,val2", list(zip(val1, val2)))
+def test_processing_fn(val1, val2):
+    # Create test data
+    df = pd.DataFrame({"val1": val1, "val2": val2})
+
+    # Get result
+    result = processing_fn(df)

data_science_tools/pandera_hypothesis/test4.py

@@ -1,26 +1,29 @@
-import hypothesis
+import pandas as pd
+import pytest
 import pandera as pa
+import hypothesis
 
 schema = pa.DataFrameSchema(
     {
         "val1": pa.Column(int, pa.Check.in_range(-2, 3)),
         "val2": pa.Column(int, pa.Check.in_range(-2, 3)),
     }
-)
+) 
 
-out_schema = schema.add_columns(
+expected = pa.DataFrameSchema(
     {
+        "val1": pa.Column(int, pa.Check.in_range(-2, 3)),
+        "val2": pa.Column(int, pa.Check.in_range(-2, 3)),
         "val3": pa.Column(float, pa.Check.in_range(-2, 3)),
-    },
+    }
 )
 
-
-@pa.check_output(out_schema)
+@pa.check_output(expected)
 def processing_fn(df):
     processed = df.assign(val3=df.val1 / df.val2)
     return processed
 
-
 @hypothesis.given(schema.strategy(size=5))
-def test_processing_fn(dataframe):
-    processing_fn(dataframe)
+def test_processing_fn(df):
+    # Get result
+    result = processing_fn(df)