Using mock library to patch a class method

Question

I'm writing unit tests, and I need to mock a method call so that in most of the cases it behaved as method itself except when argument gets a special value 'insert into'. Here is a simplified production code:

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()

which prints

Do something useful
Do something useful
Number of rows: 5

I can make a mock version myself using the following code:

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()

which prints

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

I want to use mock library to accomplish the same (I believe that this is the library included in python 3 ). How can I do that? (Python 2.6.2)

Answer 1

To be entirely clear, it's only included in python 3.3 (which I'm so happy to have learned, thank you!).

Otherwise, the pattern you could use is

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

But this would cover all cases, not just cases where you are sending 'insert into' . This might give you a hint as to where to look, but otherwise, I don't think what you're looking for is exactly possible with mock.