我正在尝试在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。
有没有人有任何其他想法尝试?
我认为这会奏效。设置全局变量可以存储数据帧。使用副作用在调用函数时保存数据帧。
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创建数据框夹具。