Python mock - 返回内部变量

Question

我正在尝试用 python 为脚本 (my_script) 编写测试。在这个脚本的中途，一个 pandas 数据框作为 csv 保存到 AWS。我不想每次运行此测试时都将文件保存到 AWS，但我确实希望能够在稍后的测试中检查此数据框的内容。是否可以覆盖 my_script 中的“save_to_csv”函数，以便此数据帧持续存在？我尝试编写一个简单的类来将数据框存储为属性:

class FakeStorage(object):
    def __init__(self):
        self.dataframe = None
    def save_to_csv(self, df, file_path):
        self.dataframe = df
        print('df fake saved')


FAKE_STORAGE = FakeStorage()
import my_script
my_script.save_to_csv = FAKE_STORAGE.save_to_csv

而且我还尝试将数据框作为全局变量保存:

def fake_save_to_csv(df, file_path):
    global DATAFRAME
    DATAFRAME = df
    print('df fake saved')

import my_script
my_script.save_to_csv = fake_save_to_csv

奇怪的是，在这两种情况下我都看到了 print 语句，因此函数似乎已被成功覆盖。但是，在测试中运行 my_script 之后，在第一种情况下我尝试访问 FAKE_STORAGE.dataframe 并且它是 None，在第二种情况下尝试访问 DATAFRAME 会抛出 NameError。

有没有人有任何其他想法可以尝试？

Answer 1

我认为这会奏效。设置可以存储数据框的全局变量。在调用函数时使用副作用来保存数据框。

import unittest
df = None. # global variable

class GlobalTestCase(unittest.TestCase):

    @unittest.mock.patch('script.save_to_csv')
    def test_aws_pull(self, mock_save_to_csv):
        mock_save_to_csv.side_effect = lambda x: df=x  # able to set global variable
        call_script()
        assert mock_save_to_csv.called  # test to see if used. you want to replace with what youd like tested

    def test_df(self):
        expected_df = 'what you expect the df to be'
        assert expected_df == df  # using thr gloval variable in the next test

其他可能的解决方案是在同一测试函数中测试数据帧，将数据帧保存为类变量而不是全局变量，将数据帧保存为单独类中的 setUp 函数，或使用 pytest 创建数据帧夹具。