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

Question

我经常需要编写一个命令行脚本来读取数据库、执行一些分析并将结果写回数据库。我努力解耦并创建一个单独的数据层通常是编写脚本 load.py、write.py 和 do_analytics.py，其中加载并编写数据库交互，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

主.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 层对于能够模拟非常重要。职责分离将避免出现意大利面条代码。