如何在Python中为模拟函数提供条件参数？

Question

我的项目使用Python的 urllib2.urlopen 来调用各种外部API。我在做单元测试时使用NoseTests，并用MiniMock来模拟对 urllib2.urlopen 的调用。

这是我的模拟代码：

from hashlib import md5
from os.path import dirname, join
from urllib2 import Request, urlopen

from minimock import mock, restore

def urlopen_stub(url, data=None, timeout=30):
    """
    Mock urllib2.urlopen and return a local file handle or create file if
    not existent and then return it.
    """

    if isinstance(url, Request):
        key = md5(url.get_full_url()).hexdigest()
    else:
        key = md5(url).hexdigest()
    data_file = join(dirname(__file__), 'cache', '%s.xml' % key)
    try:
        f = open(data_file)
    except IOError:
        restore() # restore normal function
        data = urlopen(url, data).read()
        mock('urlopen', returns_func=urlopen_stub, tracker=None) # re-mock it.
        with open(data_file, 'w') as f:
            f.write(data)
        f = open(data_file, 'r')
    return f

mock('urlopen', returns_func=urlopen_stub, tracker=None)

我这样运行我的测试：

from os.path import isfile, join
from shutil import copytree, rmtree

from nose.tools import assert_raises, assert_true

import urlopenmock

class TestMain(object):
    working = 'testing/working'

    def setUp(self):
        files = 'testing/files'
        copytree(files, self.working)

    def tearDown(self):
        rmtree(self.working)

    def test_foo(self):
        func_a_calling_urlopen()
        assert_true(isfile(join(self.working, 'file_b.txt')))

    def test_bar(self):
        func_b_calling_urlopen()
        assert_true(isfile(join(self.working, 'file_b.txt')))

    def test_bar_exception(self):
        assert_raises(AnException, func_c_calling_urlopen)

最开始，我把测试放在一个单独的模块里，这个模块导入了一个不同的模拟文件，当调用 urlopen 时返回一个损坏的XML文件。然而，导入那个模拟类后，覆盖了上面显示的模拟，导致所有测试都失败，因为每次都使用了损坏的XML。

我猜这是因为异常测试模块是在其他模块之后加载的，所以它的导入是最后执行的，而返回损坏XML的模拟函数覆盖了原来的模拟函数。

我希望能让模拟代码在运行test_bar_exception时使用损坏的XML文件，以便引发异常。我该怎么做呢？

Answer 1

假设你输入的请求是'a url'，那么输出就是'aresponse'；如果输入是'burl'，那么输出就是'bresponse'。所以可以使用

@mock.patch('requests.get', mock.Mock(side_effect = lambda k:{'aurl': 'a response', 'burl' : 'b response'}.get(k, 'unhandled request %s'%k)))

Answer 2

看起来你需要在每个需要使用模拟的urlopen的测试中设置和拆除这个模拟，并且对于那些需要处理错误情况的测试，你还需要一个不同的模拟urlopen来返回一个损坏的xml文件。就像这样：

class TestMain(object):
    # ...

    def test_foo(self):
        mock('urlopen', returns_func=urlopen_stub, tracker=None)
        func_a_calling_urlopen()
        assert_true(isfile(join(self.working, 'file_b.txt')))
        restore()

    # ...

    def test_bar_exception(self):
        mock('urlopen', 
                returns_func=urlopen_stub_which_returns_broken_xml_file, 
                tracker=None)
        assert_raises(AnException, func_c_calling_urlopen)
        restore()

不过，上面的做法有个问题，如果一个测试抛出了异常，而restore()这个调用没有被执行到，就会出错。你可以用try: ... finally:来包裹起来，但这样每个测试都要写很多额外的代码，显得有点繁琐。

你可以看看Michael Foord的mock库。

还有Pycon的演讲：http://blip.tv/file/4881513

再看看patch，它提供了装饰器和上下文管理器的选项，可以用来替换和恢复你对urlopen的调用。这个功能非常方便！