Python - 使用 side_effect 模拟在类的 init 内部调用的函数

Question

我有这个类 Foo，它的函数使用我在其构造函数中初始化的数据帧。我想在我的测试类 FooTest 中测试它的功能。

from src.shared.utils import get_spark_dataframe

class Foo(object):
    def __init__(self, x, y):
        self.a = get_spark_dataframe(x, y.some_db, "table_a")
        self.b = get_spark_dataframe(x, y.some_db, "table_b")

    def some_foo_function(self):
         return self.a.join(self.b, ['pk'])

我想模拟这个 get_spark_dataframe 函数并将其替换为我自己的函数，因为我只想用我在测试类中定义的假数据帧替换类中的数据帧。

def get_spark_dataframe(x, db_name, table_name):
    return x.get_table(db=db_name, table=table_name).toDF()

这大概是我的测试类的样子:

class FooTest(PysparkTest):
    def setUp(self):
         self.a_df = self.spark.createDataFrame([Row(...)])
         self.b_df = self.spark.createDataFrame([Row(...)])

         self.x = None
         self.y = None

    def mock_get_spark_dataframe(self, x, db_name, table_name):
         if table_name == "table_a":
               return self.a_df
         elif table_name == "table_b":
               return self.b_df

    @patch('src.shared.utils.get_spark_dataframe', side_effect=mock_get_spark_dataframe)
    def test_some_foo_function(self, mock_get_spark_dataframe):
        foo = Foo(self.x, self.y)
        return_value = foo.some_foo_function()
        ...

我做错了什么吗？当我创建 Foo 对象时，似乎没有使用我的模拟函数。真正的 get_spark_dataframe 函数似乎被调用，并且它提示 x 为 None。我是否错误地使用了 side_effect？

Answer 1

在代码中尝试这些更改:

import src.shared.utils as utils

class Foo(object):
    def __init__(self, x, y):
        self.a = utils.get_spark_dataframe(x,...
        ...

和

class FooTest(PysparkTest):
    ...
    @patch('utils.get_spark_dataframe',...
    ...