如何使用 unittest 修补具有不同 return 值的链式函数？

Question

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

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。但我无法这样做，因为 return 值是 None，但 my_function 试图在 None 值上获得 [0]['value']。

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

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

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

中将 sql_output 值设置为 None

Answer 1

编辑：我现在明白什么意思了。也许尝试用这样的测试 class 来打补丁：

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())

Answer 2

用您当前的代码解决这个问题的侵入性最小的方法是使用这样的方法：


import pyspark.sql

class SomethingTest(unittest.TestCase):

  @mock.patch.object(pyspark.sql, 'SparkSession')
  def test_my_function(self, mock_session):
    mock_session.sql.return_value.select.return_value.collect.return_value = [
        {'value': None},
    ]
    # This is the same value that thebadgateway's answer suggests.

    # the rest of your test

但是，在模拟测试中侵入性较小通常会更好。是否有一个简单的 DataFrame 可以构建，可以用于 SparkSession.sql 的 return 值？这样，您还可以确保您的 .select() 和 .collect() 回合按照您的预期进行。

这看起来像：

  @mock.patch.object(pyspark.sql, 'SparkSession')
  def test_my_function(self, mock_session):
    my_dataframe = pyspark.sql.DataFrame(...)  # build your frame
    mock_session.sql.return_value = my_dataframe

    # the rest of your test

虽然这不是直接响应您的特定问题，但使用 mock.patch.object 通常是更好的主意，因为这样可以让您直接引用要修补的对象，而不是依赖于按名称作为字符串搜索它。

如何使用 unittest 修补具有不同 return 值的链式函数？

How to patch a chained function with different return value with unittest?

python

unit-testing

python-unittest

apache-spark