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

Question

我编写了 pyspark 代码，其中涉及收集数据并转换为 CSV 操作。我可以使用以下补丁模拟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调用。

有没有办法做到这一点？

标签： pythonpython-3.xapache-spark-sqlmonkeypatchingpytest-mock