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