updated test cases

wednesday-solutions · Oct 21, 2024 · 96f7027 · 96f7027
1 parent 1f5cef7
commit 96f7027
Showing 1 changed file with 9 additions and 18 deletions.
diff --git a/tests/test_spark_wrapper_failure.py b/tests/test_spark_wrapper_failure.py
@@ -30,13 +30,11 @@ def test_value_counts_invalid_column(self):
             value_counts(self.df, "nonexistent_column")
 
         expected_error_message_1 = re.compile("Column '.+' does not exist")
-        expected_error_message_2 = re.compile("cannot resolve '.+' given input columns")
+        expected_error_message_2 = re.compile("cannot resolve '.+' given input columns:")
         actual_error_message = str(context.exception)
 
-        self.assertTrue(
-            expected_error_message_1.search(actual_error_message)
-            or expected_error_message_2.search(actual_error_message)
-        )
+        self.assertTrue(expected_error_message_1.search(actual_error_message)
+                        or expected_error_message_2.search(actual_error_message))
 
     def test_create_frame_invalid_path(self):
         with self.assertRaises(U.AnalysisException) as context:
@@ -53,29 +51,22 @@ def test_make_window_invalid_window_spec(self):
             self.df.withColumn("literal_1", F.lit(1).over(window_spec))
 
         expected_error_message_1 = re.compile("Column '.+' does not exist")
-        expected_error_message_2 = re.compile("cannot resolve '.+' given input columns")
+        expected_error_message_2 = re.compile("cannot resolve '.+' given input columns:")
         actual_error_message = str(context.exception)
 
-        self.assertTrue(
-            expected_error_message_1.search(actual_error_message)
-            or expected_error_message_2.search(actual_error_message)
-        )
+        self.assertTrue(expected_error_message_1.search(actual_error_message)
+                        or expected_error_message_2.search(actual_error_message))
 
     def test_make_window_invalid_range(self):
         with self.assertRaises(U.AnalysisException) as context:
             window_spec = make_window("market", "date", 5, 2)
             self.df.withColumn("literal_1", F.lit(1).over(window_spec))
 
         expected_error_message_1 = "The lower bound of a window frame must be less than or equal to the upper bound"
-        expected_error_message_2 = re.compile(
-            "The data type of the lower bound '.+' does not match the expected data type '.+'"
-        )
+        exoected_error_message_2 = re.compile("The data type of the lower bound '.+' does not match the expected data type '.+'")
         actual_error_message = str(context.exception)
-
-        self.assertTrue(
-            expected_error_message_1 in actual_error_message
-            or expected_error_message_2.search(actual_error_message)
-        )
+        self.assertTrue(expected_error_message_1 in actual_error_message
+                        or exoected_error_message_2.search(actual_error_message))
 
     def test_rename_column_invalid_column(self):
         with self.assertRaises(ValueError) as context: