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

时间:2017-06-01 18:23:38

标签: python unit-testing testing mocking integration-testing

我经常需要编写一个命令行脚本,该脚本将从数据库中读取,执行一些分析,并将结果写回数据库。我解耦和创建单独数据层的工作通常是编写脚本load.pywrite.pydo_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那样进行模拟,但我不确定应该嘲笑什么级别的对象,如果我有一个模拟对象,如何正确地重构实际测试,如果这是开始的最佳方法。如果之前出现这种情况,我已经在实际数据库中通过一个小的固定测试表获取解决方案,但它从来不是优雅或干净的代码。

1 个答案:

答案 0 :(得分: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

<强> main.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对象。

管理员和提供者层对于能够模拟非常重要。 单独的责任将避免使用意大利面条代码。