在 Python 中测试数据库读/写和分析脚本

Question

我经常需要编写一个命令行脚本，该脚本将从数据库中读取数据，执行一些分析，然后将结果写回数据库。我解耦和创建单独数据层的努力通常是编写脚本load.py、write.py 和do_analytics.py，其中 load 和 write 进行数据库交互，do_analytics.py 文件是这样的：

import load
import write

def batch_classify(model_filepath='my_model.pkl'):
    with open(model_filepath, 'rb') as infile:
        model = pickle.load(infile)

    data_loader = load.DataLoader()
    data_loader.load_data()
    data_loader.clean_data()
    data = data_loader.data
    # Maybe do some more manipulations here...
    output = model.transform(data)

    data_writer = write.DataWriter()
    data_writer.write_data(output)

if __name__ == "__main__":
    # maybe would have some command line options here to pass to batch_classify
    batch_classify()   

我现在想测试一些固定的数据集，并确保分类（输出）结果符合我的预期。我现在不需要测试实际的数据库连接，所以根据一些研究，我想我想像this post 那样进行模拟，但我不确定应该模拟什么级别的对象，如何正确重构为实际测试一旦我有了模拟对象，如果这甚至是最好的方法。之前出现这种情况时，我四处寻找可以通过实际数据库中的小型固定测试表工作的解决方案，但它从来都不是优雅或干净的代码。

Answer 1

在你的情况下，我会有 4 个文件。

database_provider.py

class DatabaseProvider(object):
    def get_data(self):
        return db.get() # Get your data

    def set_data(self, data):
        db.set(data) # update your data

analytic_manager.py

class AnalyticManager(object):
    def __init__(self):
        self.database_provider = DatabaseProvider()

    def process(self, arguments):
        # Get data from DB
        data = self.database_provider.get_data()

        # Do your logic here
        data = self.clean(data)
        data = self.transform(data)

        # Save in DB
        self.database_provider.set_data(data)

    def clean(self, data):
        # do cleaning
        return cleaned_data

    def transform(self, data):
        # do transform
        return transformed_data

ma​​in.py

if __name__ == "__main__":
    arguments = whatever
    manager = AnalyticManager(arguments)
    manager.process(arguments)

test_analytic_manager.py

import unittest
from mock import Mock, patch

class TestAnalyticManager(unittest.TestCase):


    @patch("database_provider.DatabaseProvider.get_data")
    @patch("database_provider.DatabaseProvider.set_data")
    def test_process_should_clean_and_transform_data(self, mock_set_data, mock_get_data):
        # Arranges
        arguments = whatever
        manager = AnalyticManager(arguments)

        mock_get_data.return_value = ["data from DB", "data2 from DB"]

        expected_data = ["cleaned and transformed data1", "cleaned and transformed data2"]

        # Acts
        manager.process(arguments)

        # Asserts
        mock_set_data.assert_called_once_with(expected_data)

您现在可以根据需要模拟您的提供程序。 最重要的是在管理器内部而不是提供者内部执行所有逻辑。 您的 db_provider 应该只与您的数据库进行交互并将接收到的数据映射到您的 python 对象。

Manager 和 Provider 层对于能够模拟非常重要。 单独的职责将避免出现意大利面条式代码。