refactor interpol_tests complete

databrickslabs · Jul 9, 2024 · 09e6423 · 09e6423
1 parent ab5210a
commit 09e6423
Show file tree

Hide file tree

Showing 2 changed files with 1,524 additions and 1,493 deletions.
diff --git a/python/tests/interpol_tests.py b/python/tests/interpol_tests.py
@@ -24,7 +24,7 @@ def test_validate_fill_method(self):
         )
 
     def test_validate_col_exist_in_df(self):
-        input_df: DataFrame = self.get_data_as_sdf("input_data")
+        input_df: DataFrame = self.get_test_df_builder("init").as_sdf()
 
         self.assertRaises(
             ValueError,
@@ -54,7 +54,7 @@ def test_validate_col_exist_in_df(self):
         )
 
     def test_validate_col_target_cols_data_type(self):
-        input_df: DataFrame = self.get_data_as_sdf("input_data")
+        input_df: DataFrame = self.get_test_df_builder("init").as_sdf()
 
         self.assertRaises(
             TypeError,
@@ -69,7 +69,7 @@ def test_fill_validation(self):
         """Test fill parameter is valid."""
 
         # load test data
-        input_tsdf: TSDF = self.get_data_as_tsdf("input_data")
+        input_tsdf: TSDF = self.get_test_df_builder("init").as_tsdf()
 
         # interpolate
         self.assertRaises(
@@ -89,7 +89,7 @@ def test_target_column_validation(self):
         """Test target columns exist in schema, and are of the right type (numeric)."""
 
         # load test data
-        input_tsdf: TSDF = self.get_data_as_tsdf("input_data")
+        input_tsdf: TSDF = self.get_test_df_builder("init").as_tsdf()
 
         # interpolate
         self.assertRaises(
@@ -109,7 +109,7 @@ def test_partition_column_validation(self):
         """Test partition columns exist in schema."""
 
         # load test data
-        input_tsdf: TSDF = self.get_data_as_tsdf("input_data")
+        input_tsdf: TSDF = self.get_test_df_builder("init").as_tsdf()
 
         # interpolate
         self.assertRaises(
@@ -129,7 +129,7 @@ def test_ts_column_validation(self):
         """Test time series column exist in schema."""
 
         # load test data
-        input_tsdf: TSDF = self.get_data_as_tsdf("input_data")
+        input_tsdf: TSDF = self.get_test_df_builder("init").as_tsdf()
 
         # interpolate
         self.assertRaises(
@@ -154,8 +154,8 @@ def test_zero_fill_interpolation(self):
         """
 
         # load test data
-        simple_input_tsdf: TSDF = self.get_data_as_tsdf("simple_input_data")
-        expected_df: DataFrame = self.get_data_as_sdf("expected_data")
+        simple_input_tsdf: TSDF = self.get_test_df_builder("simple_init").as_tsdf()
+        expected_df: DataFrame = self.get_test_df_builder("expected").as_sdf()
 
         # interpolate
         actual_df: DataFrame = self.interpolate_helper.interpolate(
@@ -180,8 +180,8 @@ def test_zero_fill_interpolation_no_perform_checks(self):
         """
 
         # load test data
-        simple_input_tsdf: TSDF = self.get_data_as_tsdf("simple_input_data")
-        expected_df: DataFrame = self.get_data_as_sdf("expected_data")
+        simple_input_tsdf: TSDF = self.get_test_df_builder("simple_init").as_tsdf()
+        expected_df: DataFrame = self.get_test_df_builder("expected").as_sdf()
 
         # interpolate
         actual_df: DataFrame = self.interpolate_helper.interpolate(
@@ -207,8 +207,8 @@ def test_null_fill_interpolation(self):
         """
 
         # load test data
-        simple_input_tsdf: TSDF = self.get_data_as_tsdf("simple_input_data")
-        expected_df: DataFrame = self.get_data_as_sdf("expected_data")
+        simple_input_tsdf: TSDF = self.get_test_df_builder("simple_init").as_tsdf()
+        expected_df: DataFrame = self.get_test_df_builder("expected").as_sdf()
 
         # interpolate
         actual_df: DataFrame = self.interpolate_helper.interpolate(
@@ -234,8 +234,8 @@ def test_back_fill_interpolation(self):
         """
 
         # load test data
-        simple_input_tsdf: TSDF = self.get_data_as_tsdf("simple_input_data")
-        expected_df: DataFrame = self.get_data_as_sdf("expected_data")
+        simple_input_tsdf: TSDF = self.get_test_df_builder("simple_init").as_tsdf()
+        expected_df: DataFrame = self.get_test_df_builder("expected").as_sdf()
 
         # interpolate
         actual_df: DataFrame = self.interpolate_helper.interpolate(
@@ -261,8 +261,8 @@ def test_forward_fill_interpolation(self):
         """
 
         # load test data
-        simple_input_tsdf: TSDF = self.get_data_as_tsdf("simple_input_data")
-        expected_df: DataFrame = self.get_data_as_sdf("expected_data")
+        simple_input_tsdf: TSDF = self.get_test_df_builder("simple_init").as_tsdf()
+        expected_df: DataFrame = self.get_test_df_builder("expected").as_sdf()
 
         # interpolate
         actual_df: DataFrame = self.interpolate_helper.interpolate(
@@ -288,8 +288,8 @@ def test_linear_fill_interpolation(self):
         """
 
         # load test data
-        simple_input_tsdf: TSDF = self.get_data_as_tsdf("simple_input_data")
-        expected_df: DataFrame = self.get_data_as_sdf("expected_data")
+        simple_input_tsdf: TSDF = self.get_test_df_builder("simple_init").as_tsdf()
+        expected_df: DataFrame = self.get_test_df_builder("expected").as_sdf()
 
         # interpolate
         actual_df: DataFrame = self.interpolate_helper.interpolate(
@@ -313,8 +313,8 @@ def test_different_freq_abbreviations(self):
         """
 
         # load test data
-        simple_input_tsdf: TSDF = self.get_data_as_tsdf("simple_input_data")
-        expected_df: DataFrame = self.get_data_as_sdf("expected_data")
+        simple_input_tsdf: TSDF = self.get_test_df_builder("simple_init").as_tsdf()
+        expected_df: DataFrame = self.get_test_df_builder("expected").as_sdf()
 
         # interpolate
         actual_df: DataFrame = self.interpolate_helper.interpolate(
@@ -340,8 +340,8 @@ def test_show_interpolated(self):
         """
 
         # load test data
-        simple_input_tsdf: TSDF = self.get_data_as_tsdf("simple_input_data")
-        expected_df: DataFrame = self.get_data_as_sdf("expected_data")
+        simple_input_tsdf: TSDF = self.get_test_df_builder("simple_init").as_tsdf()
+        expected_df: DataFrame = self.get_test_df_builder("expected").as_sdf()
 
         # interpolate
         actual_df: DataFrame = self.interpolate_helper.interpolate(
@@ -358,7 +358,7 @@ def test_show_interpolated(self):
         self.assertDataFrameEquality(expected_df, actual_df, ignore_nullable=True)
 
     def test_validate_ts_col_data_type_is_not_timestamp(self):
-        input_df: DataFrame = self.get_data_as_sdf("input_data")
+        input_df: DataFrame = self.get_test_df_builder("init").as_sdf()
 
         self.assertRaises(
             ValueError,
@@ -374,7 +374,7 @@ def test_interpolation_freq_is_none(self):
         """Test a ValueError is raised when freq is None."""
 
         # load test data
-        simple_input_tsdf: TSDF = self.get_data_as_tsdf("input_data")
+        simple_input_tsdf: TSDF = self.get_test_df_builder("init").as_tsdf()
 
         # interpolate
         self.assertRaises(
@@ -394,7 +394,7 @@ def test_interpolation_func_is_none(self):
         """Test a ValueError is raised when func is None."""
 
         # load test data
-        simple_input_tsdf: TSDF = self.get_data_as_tsdf("input_data")
+        simple_input_tsdf: TSDF = self.get_test_df_builder("init").as_tsdf()
 
         # interpolate
         self.assertRaises(
@@ -414,7 +414,7 @@ def test_interpolation_func_is_callable(self):
         """Test ValueError is raised when func is callable."""
 
         # load test data
-        simple_input_tsdf: TSDF = self.get_data_as_tsdf("input_data")
+        simple_input_tsdf: TSDF = self.get_test_df_builder("init").as_tsdf()
 
         # interpolate
         self.assertRaises(
@@ -434,7 +434,7 @@ def test_interpolation_freq_is_not_supported_type(self):
         """Test ValueError is raised when func is callable."""
 
         # load test data
-        simple_input_tsdf: TSDF = self.get_data_as_tsdf("input_data")
+        simple_input_tsdf: TSDF = self.get_test_df_builder("init").as_tsdf()
 
         # interpolate
         self.assertRaises(
@@ -459,8 +459,8 @@ def test_interpolation_using_default_tsdf_params(self):
         """
 
         # load test data
-        simple_input_tsdf: TSDF = self.get_data_as_tsdf("simple_input_data")
-        expected_df: DataFrame = self.get_data_as_sdf("expected")
+        simple_input_tsdf: TSDF = self.get_test_df_builder("simple_init").as_tsdf()
+        expected_df: DataFrame = self.get_test_df_builder("expected").as_sdf()
 
         # interpolate
         actual_df: DataFrame = simple_input_tsdf.interpolate(
@@ -475,8 +475,8 @@ def test_interpolation_using_custom_params(self):
         modified params."""
 
         # Modify input DataFrame using different ts_col
-        simple_input_tsdf: TSDF = self.get_data_as_tsdf("simple_input_data")
-        expected_df: DataFrame = self.get_data_as_sdf("expected")
+        simple_input_tsdf: TSDF = self.get_test_df_builder("simple_init").as_tsdf()
+        expected_df: DataFrame = self.get_test_df_builder("expected").as_sdf()
 
         input_tsdf = TSDF(
             simple_input_tsdf.df.withColumnRenamed("event_ts", "other_ts_col"),
@@ -501,7 +501,7 @@ def test_tsdf_constructor_params_are_updated(self):
         interpolation."""
 
         # load test data
-        simple_input_tsdf: TSDF = self.get_data_as_tsdf("simple_input_data")
+        simple_input_tsdf: TSDF = self.get_test_df_builder("simple_init").as_tsdf()
 
         actual_tsdf: TSDF = simple_input_tsdf.interpolate(
             ts_col="event_ts",
@@ -520,8 +520,8 @@ def test_interpolation_on_sampled_data(self):
         """Verify interpolation can be chained with resample within the TSDF class"""
 
         # load test data
-        simple_input_tsdf: TSDF = self.get_data_as_tsdf("simple_input_data")
-        expected_df: DataFrame = self.get_data_as_sdf("expected")
+        simple_input_tsdf: TSDF = self.get_test_df_builder("simple_init").as_tsdf()
+        expected_df: DataFrame = self.get_test_df_builder("expected").as_sdf()
 
         actual_df: DataFrame = (
             simple_input_tsdf.resample(freq="30 seconds", func="mean", fill=None)
@@ -538,8 +538,8 @@ def test_defaults_with_resampled_df(self):
         # self.buildTestingDataFrame()
 
         # load test data
-        simple_input_tsdf = self.get_data_as_tsdf("simple_input_data")
-        expected_df: DataFrame = self.get_data_as_sdf("expected", convert_ts_col=True)
+        simple_input_tsdf = self.get_test_df_builder("simple_init").as_tsdf()
+        expected_df: DataFrame = self.get_test_df_builder("expected").as_sdf()
 
         actual_df: DataFrame = (
             simple_input_tsdf.resample(freq="30 seconds", func="mean", fill=None)