利用Python Mock模拟OJ test case

场景需要

在刷OJ的过程中，最常见的就是从stdin里面读取输入了，然后把结果输出到stdout上供OJ判断结果。

一般来说，在本地写完程序之后都会手动输入一遍test case，观察输出结果之后发现不对劲，再手动输入一遍test case。。。

作为一个搞笑的程序员，啊呸，不是，高效的程序员，简直不能忍受一次又一次地手动输入test case，能不能每次debug完之后一键跑数据呢？而且能不能让自己额外想出来的test case保存下来，每次都跑一遍呢？

由于我现在主要是使用Python刷OJ，因此很自然地想到了通过Python Mock来模拟stdin。

例子

以我最近的刷的PAT 1002 A+B for Polynomials (25) Python为例，要完整的模拟整个OJ测试过程主要从两个方面入手，输入和输出。

模拟输入

在Python从console读取输入主要是通过input函数，那么通过一个Mock对象代替掉input函数即可。但是input函数是一个内置函数，怎么Mock呢？

在Python3中，如果要标记内置函数，可以通过builtins包导入；在Python2中，则是通过__builtin__包导入。

Python 3.6.1 (v3.6.1:69c0db5050, Mar 21 2017, 01:21:04)[GCC 4.2.1 (Apple Inc. build 5666) (dot 3)] on darwinType "help", "copyright", "credits" or "license" for more information.>>> import builtins>>> builtins.input<built-in function input>Python 2.7.11 (default, Jan 22 2016, 08:29:18)[GCC 4.2.1 Compatible Apple LLVM 7.0.2 (clang-700.1.81)] on darwinType "help", "copyright", "credits" or "license" for more information.>>> import __builtin__>>> __builtin__.input<built-in function input>

stackoverflow传送门：Where is the builtin module in Python3? Why was it renamed?

接下来只要在测试函数上加上patch装饰器即可：

class Test1002(object):    @patch('builtins.input', lambda: '2 1 2.4 0 3.2')    def test_1(self):        assert input() == '2 1 2.4 0 3.2'

在上面的代码例子里面我使用一个lambda函数替换了内置函数input，这样就能模拟输入的数据了。

但是还有一个问题，就是运行的程序可能不止调用input一次，可能调用多次，要求每次调用input返回的数据都不一样，这时可以通过一个具有side_effect的Mock类进行模拟。什么是side_effect呢？通过一个实例来说明：

Python 3.6.1 (v3.6.1:69c0db5050, Mar 21 2017, 01:21:04)[GCC 4.2.1 (Apple Inc. build 5666) (dot 3)] on darwinType "help", "copyright", "credits" or "license" for more information.>>> from mock import Mock>>> m = Mock(side_effect=['1', '2', '3', '4'])>>> m()'1'>>> m()'2'>>> m()'3'

从例子上来看，就是每次调用具有side_effect的Mock对象都会根据list的顺序返回不同的对象，用来替代input函数的多次输入刚好。

class Test1002(object):    @patch('builtins.input', Mock(side_effect=['2 1 2.4 0 3.2', '2 2 1.5 1 0.5']))    def test_2(self):        assert input() == '2 1 2.4 0 3.2'        assert input() == '2 2 1.5 1 0.5'

捕获输出

OJ基本就是通过stdout输出的数据进行判定，因此如果能把stdout的数据进行重定向即可拿到输出数据。还有另外一种方法是不重定向stdout，但是在stdout之前把数据记录下来。

参考了fluent python中的一个反转sys.stdout.write的例子，我通过一个上下文管理器实现了这个功能。(在离开上下文的时候回复原来的sys.stdout.write)

class PrintRewrite(object):    def __init__(self):        self.res = ''    def __enter__(self):        import sys        self.origin_write = sys.stdout.write        sys.stdout.write = self.log_print        return self    def log_print(self, text):        self.res += text        self.origin_write(text)    def __exit__(self, exc_type, exc_val, exc_tb):        import sys        sys.stdout.write = self.origin_write

在每次print之前把print的结果添加到res变量中，因此，如果要检验测试结果：

class Test1002(object):    @patch('builtins.input', lambda: '2 1 2.4 0 3.2')    def test_1(self):        with PrintRewrite() as t:            main()            assert t.res == '2 1 4.8 0 6.4'    @patch('builtins.input', Mock(side_effect=['2 1 2.4 0 3.2', '2 2 1.5 1 0.5']))    def test_2(self):        with PrintRewrite() as t:            main()            assert t.res == '3 2 1.5 1 2.9 0 3.2'

需要注意的是在Python3中，sys.stdout.write方法才可以被重写。在Python2中，write属性是一个只读属性，sys.stdout需要通过一个实现了类似功能的file对象替代才行。
stackoverflow传送门：Set a Read-Only Attribute in Python?