Add exception handler

riyaverm-db · May 29, 2024 · 08350ad · 08350ad
1 parent 2a50ba1
commit 08350ad
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python/pyspark/sql/tests/test_python_streaming_datasource.py b/python/pyspark/sql/tests/test_python_streaming_datasource.py
@@ -27,6 +27,7 @@
     SimpleDataSourceStreamReader,
     WriterCommitMessage,
 )
+from pyspark.sql.streaming import StreamingQueryException
 from pyspark.sql.types import Row
 from pyspark.testing.sqlutils import (
     have_pyarrow,
@@ -231,6 +232,8 @@ def test_stream_writer(self):
                 [Row("failed in batch 1")],
             )
             q.awaitTermination()
+        except StreamingQueryException as e:
+            self.assertIn("invalid value", str(e))
         finally:
             input_dir.cleanup()
             output_dir.cleanup()