From 08350adaa6a2dcaec0bc7de656569d6b54552a8e Mon Sep 17 00:00:00 2001 From: Riya Verma Date: Wed, 29 May 2024 11:10:54 -0700 Subject: [PATCH] Add exception handler --- python/pyspark/sql/tests/test_python_streaming_datasource.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/pyspark/sql/tests/test_python_streaming_datasource.py b/python/pyspark/sql/tests/test_python_streaming_datasource.py index 86fb49ece0283..e5622e28f15bf 100644 --- 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()