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',...
    ...

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

1 个答案: