时间戳

首页 > 解决方案 > 如何使用 unittest 修补具有不同返回值的链式函数？

python - 如何使用 unittest 修补具有不同返回值的链式函数？

问题描述

我有一个看起来像这样的函数：

def my_function():
   sql_output = spark.sql('query').select('value').collect()[0]['value']

我正在尝试在 Unittest 中使用 Mock 和 Patch 来修补变量sql_output。我正在修补spark.sql功能：

@patch("my_function.spark.sql")
def test_my_function(self, mock_sql_functions):
    from pyspark.sql.types import StringType
    from pyspark.sql.functions import lit

    mock_sql_functions.return_value.select.return_value.collect.return_value = None

我的目标是设置sql_output等于无。但我不能这样做，因为返回值为 None，但my_function试图获取[0]['value']该None值。

我尝试将返回值作为数据框，如下所示：

sdf = spark.createDataFrame([('None', 'None', 'None')], ['value', 'value2', 'value3'])
sdf = sdf.withColumn("value", lit(None).cast(StringType()))

mock_sql_functions.return_value.select.return_value.collect.return_value = sdf

但它不起作用，因为我需要使用[0]['value'], 同时collect()我相信。所以我的问题是，如何将这些倍数设置return_value为不同的值？或者我该如何设置sql_output值？Noneunittest

标签： pythonunit-testingapache-sparkpython-unittest

解决方案

编辑：我明白了现在的意思。也许尝试使用这样的测试类进行修补：

class TestSpark:
    def sql(self, arg): pass
    def select(self, arg): pass
    def collect(self): return [{"value": None}]

那么装饰器将类似于@mock.patch.object(the_module, "spark", return_value=TestSpark())

推荐阅读