如何对 Python pyspark 的实用程序函数进行单元测试

如何解决如何对 Python pyspark 的实用程序函数进行单元测试

我有一个用 Python 编写的实用程序函数，用于将 parquet 文件和 json 文件写入 s3 存储桶。 这是功能：

def write_to_s3(data1,data2,s3_path):

    try:
        data1.write.mode("overwrite").parquet(s3_path)
        data2.write.mode("overwrite").json(s3_path,compression="gzip")

    except Exception as err:
        logging.error(err)
        raise

我仍在学习单元测试，只是想知道是否有办法模拟 spark 会话以避免在单元测试中设置真正的会话？有人可以帮我为此编写单元测试用例吗？我发现了一个类似的问题，但它是针对 Scala 的，它需要设置一个 Spark 会话，我认为有一种方法可以模拟它，就像我们可以模拟 s3 一样？希望这是有道理的，谢谢。

更新：我遵循了@Mauro Baraldi 在下面推荐的page，该方法有效，但它只查看被称为第一个的 write 操作，我如何测试 {{1} }} & parquet 部分以确保数据以预期格式写入 s3？谢谢。