Python mock - вернуть внутреннюю переменную

Question

Я пытаюсь написать тест для скрипта (my_script) на python.На полпути этого скрипта кадр данных pandas сохраняется в AWS как csv.Я не хочу сохранять файл в AWS каждый раз, когда я запускаю этот тест, но я хочу иметь возможность проверить содержимое этого фрейма позже в этом тесте.Можно ли переписать функцию «save_to_csv» внутри my_script, чтобы этот кадр данных сохранялся?Я попытался написать простой класс для хранения фрейма данных как атрибута:

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, и он отсутствует, во втором случае при попытке доступа к 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 создаст объект данных.