src/tools/telemetry/telemetry/unittest/simple_mock.py

   1 # Copyright 2012 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4 """A very very simple mock object harness."""
   5 from types import ModuleType
   6
   7 DONT_CARE = ''
   8
   9 class MockFunctionCall(object):
  10   def __init__(self, name):
  11     self.name = name
  12     self.args = tuple()
  13     self.return_value = None
  14     self.when_called_handlers = []
  15
  16   def WithArgs(self, *args):
  17     self.args = args
  18     return self
  19
  20   def WillReturn(self, value):
  21     self.return_value = value
  22     return self
  23
  24   def WhenCalled(self, handler):
  25     self.when_called_handlers.append(handler)
  26
  27   def VerifyEquals(self, got):
  28     if self.name != got.name:
  29       raise Exception('Self %s, got %s' % (repr(self), repr(got)))
  30     if len(self.args) != len(got.args):
  31       raise Exception('Self %s, got %s' % (repr(self), repr(got)))
  32     for i in range(len(self.args)):
  33       self_a = self.args[i]
  34       got_a = got.args[i]
  35       if self_a == DONT_CARE:
  36         continue
  37       if self_a != got_a:
  38         raise Exception('Self %s, got %s' % (repr(self), repr(got)))
  39
  40   def __repr__(self):
  41     def arg_to_text(a):
  42       if a == DONT_CARE:
  43         return '_'
  44       return repr(a)
  45     args_text = ', '.join([arg_to_text(a) for a in self.args])
  46     if self.return_value in (None, DONT_CARE):
  47       return '%s(%s)' % (self.name, args_text)
  48     return '%s(%s)->%s' % (self.name, args_text, repr(self.return_value))
  49
  50 class MockTrace(object):
  51   def __init__(self):
  52     self.expected_calls = []
  53     self.next_call_index = 0
  54
  55 class MockObject(object):
  56   def __init__(self, parent_mock = None):
  57     if parent_mock:
  58       self._trace = parent_mock._trace # pylint: disable=W0212
  59     else:
  60       self._trace = MockTrace()
  61
  62   def __setattr__(self, name, value):
  63     if (not hasattr(self, '_trace') or
  64         hasattr(value, 'is_hook')):
  65       object.__setattr__(self, name, value)
  66       return
  67     assert isinstance(value, MockObject)
  68     object.__setattr__(self, name, value)
  69
  70   def SetAttribute(self, name, value):
  71     setattr(self, name, value)
  72
  73   def ExpectCall(self, func_name, *args):
  74     assert self._trace.next_call_index == 0
  75     if not hasattr(self, func_name):
  76       self._install_hook(func_name)
  77
  78     call = MockFunctionCall(func_name)
  79     self._trace.expected_calls.append(call)
  80     call.WithArgs(*args)
  81     return call
  82
  83   def _install_hook(self, func_name):
  84     def handler(*args, **_):
  85       got_call = MockFunctionCall(
  86         func_name).WithArgs(*args).WillReturn(DONT_CARE)
  87       if self._trace.next_call_index >= len(self._trace.expected_calls):
  88         raise Exception(
  89           'Call to %s was not expected, at end of programmed trace.' %
  90           repr(got_call))
  91       expected_call = self._trace.expected_calls[
  92         self._trace.next_call_index]
  93       expected_call.VerifyEquals(got_call)
  94       self._trace.next_call_index += 1
  95       for h in expected_call.when_called_handlers:
  96         h(*args)
  97       return expected_call.return_value
  98     handler.is_hook = True
  99     setattr(self, func_name, handler)
 100
 101
 102 class MockTimer(object):
 103   """ A mock timer to fake out the timing for a module.
 104     Args:
 105       module: module to fake out the time
 106   """
 107   def __init__(self, module=None):
 108     self._elapsed_time = 0
 109     self._module = module
 110     self._actual_time = None
 111     if module:
 112       assert isinstance(module, ModuleType)
 113       self._actual_time = module.time
 114       self._module.time = self
 115
 116   def sleep(self, time):
 117     self._elapsed_time += time
 118
 119   def time(self):
 120     return self._elapsed_time
 121
 122   def SetTime(self, time):
 123     self._elapsed_time = time
 124
 125   def __del__(self):
 126     self.Restore()
 127
 128   def Restore(self):
 129     if self._module:
 130       self._module.time = self._actual_time
 131       self._module = None
 132       self._actual_time = None