如何模拟 returns 不同对象的链式 python 函数

Question

我正在尝试模拟具有以下代码的函数 write_tables_to_datalake_by_partition_datehour()。

tables_with_data[table].repartition("datehour_partition").write.format(
                "parquet"
            ).partitionBy("datehour_partition").mode("overwrite").save(
                cleansed_file_path + table
            )

tables_with_data[table] 属于 pyspark.sql.DataFrame 类型。查看链式函数时，它保持 DataFrame 直到 write 函数 returns 变成 DataFrameWriter。 DataFrameWriter 实际上可以使用其余的功能。

我不想将数据写入我们的数据湖。我想嘲笑它。

@mock.patch("module_name.io.DataFrame")
def test_write_tables_to_datalake_by_partition_datehour(mock_df) -> None:

    mock_df.return_value.repartition.return_value.write.format.return_value.partitionBy.return_value.mode.return_value.save.return_value = (
        "yes!"
    )

这不起作用，因为 DataFrame 没有方法 save，这仅在 DataFrameWriter 中可用。但是我无法修补 DataFrameWriter 因为我没有在我的模块中导入它，只有 DataFrame 是。

我也试过了

    with mock.patch.object(DataFrameWriter, "save") as mock1:
        mock1.return_value.format.partitionBy.mode.save.return_value = "test"

        result = write_tables_to_datalake_by_partition_datehour(
            tablenames,
            dataframes_cleansed,
            cleansed_file_path,
            dataframes_cleansed,
            quarantine_file_path,
        )

这也不行。我怎样才能以正确的方式模拟它？

Answer 1

我通过在我的模块中导入 DataFrameWriter 使其工作。然后我模拟 save return 值，现在它不会写入我们的数据湖。非常简单的解决方案。

@mock.patch("mymodule.io.DataFrameWriter.save", return_value="1")
def test_write_tables_to_datalake_by_partition_datehour(mock_dfw) -> None:
   # implementation

在我的模块中，我在顶部添加了它，以便我的 mock.patch 可以找到实际的 class

from pyspark.sql import DataFrameWriter

即使它是未使用的导入，模拟也需要它。

如何模拟 returns 不同对象的链式 python 函数

How to mock a chained python function that returns different objects

mocking

dataframe

pytest

python-3.x

pyspark