test

Author: bubriks

python/tests/engine/test_spark.py

@@ -4559,8 +4559,9 @@ def test_save_empty_dataframe(self, mocker):
         mock_spark_engine_save_dataframe = mocker.patch(
             "hsfs.engine.spark.Engine.save_dataframe"
         )
-        mock_spark_format = mocker.Mock()
-        mocker.patch("pyspark.sql.session.SparkSession.read", return_value=mock_spark_format)
+        mock_spark_read = mocker.patch("pyspark.sql.session.SparkSession.read")
+        mock_spark_read_format = mock_spark_read.return_value.format
+        mock_spark_read_format.return_value = mocker.Mock()
 
         # Arrange
         spark_engine = spark.Engine()
@@ -4580,7 +4581,7 @@ def test_save_empty_dataframe(self, mocker):
 
         # Assert
         assert mock_spark_engine_save_dataframe.call_count == 1
-        assert mock_spark_format.call_count == 1
+        assert mock_spark_read_format.call_count == 1
 
     def test_apply_transformation_function_single_output(self, mocker):
         # Arrange