使用mock库修补类方法

Question

我正在写单元测试，需要模拟一个方法的调用。大部分情况下，这个方法的行为和原来的方法一样，只有当参数是特殊值 'insert into' 时，才会有不同的表现。下面是简化后的生产代码：

class CommandServer(object):
    def __init__(self):
        self.rowcount = None
    def runSQL(self, sql):
        print "Do something useful"
        self.rowcount=5
        return self

class Process(object):
    def process(self):
        cs = CommandServer()
        cs.runSQL("create table tbl1(X VARCHAR2(10))")
        r = cs.runSQL("insert into tbl1 select * from tbl2")
        print "Number of rows: %s" % r.rowcount

p = Process()
p.process()

这段代码的输出是：

Do something useful
Do something useful
Number of rows: 5

我可以自己用以下代码制作一个模拟版本：

runSQL = CommandServer.runSQL
def runSQLPatch(self, sql):
    if sql.lstrip().startswith('insert into'):
        print "Patched version in use"
        class res(object):
            rowcount = -1
        return res
    else:
        return runSQL(self, sql)
CommandServer.runSQL = runSQLPatch

p = Process()
p.process()

这段代码的输出是：

Do something useful
Patched version in use
Number of rows: -1

我想使用 mock 库 来实现同样的功能（我相信这个库包含在 python 3 中）。我该怎么做呢？（Python 2.6.2）

Answer 1

为了让大家更明白，这个功能只在 Python 3.3 版本中有（我很高兴能学到这一点，谢谢！）。

如果没有这个功能，你可以使用下面的模式：

from mock import patch

with patch.object(CommandServer, 'runSQL') as runSQL:
    class res(object):
       rowcount = -1

    runSQL.return_value = res

    p = Process()
    p.process()

    for c in runSQL.call_list:
        assert c[1].lstrip().startswith('insert into') is True

不过，这种方法会涵盖所有情况，而不仅仅是你在发送 'insert into' 的时候。这可能会给你一些线索，让你知道该往哪里找，但除此之外，我觉得你想要的功能用 mock 是实现不了的。