pytest monkeypatch：如何模拟 spark.sql("query").coalesce(1).write.csv(path, header=True)

如何解决pytest monkeypatch：如何模拟 spark.sql("query").coalesce(1).write.csv(path, header=True)

我编写了涉及收集数据和转换为 CSV 操作的 pyspark 代码。我可以使用以下补丁模拟 collect() 的 spark.sql 方法：

class MockSpark:
    @staticmethod
    def collect(*args,**kwargs):
        return []

def mock_return(*args,**kwargs):
    return MockSpark(*args,**kwargs)

monkeypatch.setattr(SparkSession,"sql",mock_return)

但是，我无法模拟下面的语句，因为它涉及调用 write.csv()：

spark.sql("select * from temp_table").coalesce(1).write.csv(path,header=True)

由于调用 write.csv()，我无法以类似于 coalesce 调用的方式模拟 collect。

有没有办法做到这一点？