A python mock object factory

oremock.py

#
# written by ooharak 2012. https://github.com/ooharak
# Usage:
#   # an example test target 
#   def target():
#     import os
#     os.putenv('DUMMY', os.getenv('PATH')+'#'+os.getenv('USER'))
#
#   ...
#
#   # setup mock
#   import oremock
#   m = oremock.OreMock()
#   m.to_mock('os.getenv', ['/bin','scott'])
#   m.to_mock('os.putenv')
#
#   # call the test target
#   try:
#     target() 
#
#     # check phase
#     assertEqual(['os.getenv', 'os.getenv', 'os.putenv'], m.callorder)
#     assertEqual('PATH', m.getcall('os.getenv')[0])
#     assertEqual('USER', m.getcall('os.getenv')[0])
#     assertEqual(('DUMMY','/bin#scott'), m.getcall('os.putenv'))
#     # check that no more calls were made
#     assertFalse(m.hasCalled('os.putenv'))
#     assertFalse(m.hasCalled('os.getenv'))
#   finally:
#     m.unmock_all()
#   
#   
#   
#   
#
class OreMock(object):
  """tiny mock object factory"""

  def __init__(self):
    self.callorder = []
    self.calls = {}
    self.orgs = {}
    self.objnames = {}
    self.rets = {}

  def _mock(self, name, *param, **nparam):
   import exceptions
   #import java.lang
   pdata = param
   if len(nparam) > 0:
     pdata = param + (nparam,)

   self.callorder.append(name)
   if self.calls.has_key(name):
     self.calls[name].append(pdata)
   else:
     self.calls[name] = [pdata]
   if self.rets[name] == None:
     return None
   else:
     try:
       ret = self.rets[name].pop(0)
       if isinstance(ret, exceptions.Exception):
         raise ret
       #elif isinstance(ret, java.lang.Exception):
       # raise ret
       else:
         return ret
     except IndexError:
       raise 'No more call of ' + name

  # shift the parameter by one
  def getcall(self, name):
    try:
      return self.calls[name].pop(0)
    except KeyError:
      raise 'Never called: ' + name
    except IndexError:
      raise 'No more call of '+name

  # returns whether this mock has one or more calls recorded but yet not getcall()'d
  def hasCalled(self, name):
    return self.calls.has_key(name) and len(self.calls[name]) > 0

  # replace a function of the name objname to a mock.
  # each call to that mock will return an element of the array specified with ret.
  # name is an alias for the mock.
  def to_mock_as(self, name, objname, ret=None):
    self.objnames[name] = objname
    if '.' in objname:
      exec('import ' + '.'.join(objname.split('.')[0:-1])) in globals(), locals()
    exec('orgval = '+objname) in globals(), locals()
    self.orgs[name] = orgval
    self.rets[name] = ret
    m = lambda *p, **pp: self._mock(name, *p, **pp)
    exec(objname + ' = m') in globals(), locals()

  def to_mock(self, objname, ret=None):
    return self.to_mock_as(objname, objname, ret)
  
  def unmock_all(self):
    for i in self.objnames.items():
      name = i[0]
      objname = i[1]
      if '.' in objname:
        exec('import ' + '.'.join(objname.split('.')[0:-1])) in globals(), locals()
      obj = self.orgs[name]
      exec(objname + ' = obj') in globals(), locals()