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

Question

I have this class, Foo, whose functions use dataframes I initialize in its constructor. I want to test its functions in my test class, 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'])

I want to mock out this get_spark_dataframe function and replace it with my own, since I'm merely interested in replacing the dataframes in the class with fake dataframes I define in my test class.

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

This is vaguely what my test class looks like:

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

Am I doing something incorrectly? It doesn't seem like my mock function is being used when I create the Foo object. The real get_spark_dataframe function seems to be getting called, and it complains about x being None. Am I using side_effect incorrectly?

Answer 1

Try these changes in your code:

import src.shared.utils as utils

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

and

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