1fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville#!/usr/bin/python2.4
2fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville#
3fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# Copyright 2008 Google Inc.
4fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville#
5fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# Licensed under the Apache License, Version 2.0 (the "License");
6fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# you may not use this file except in compliance with the License.
7fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# You may obtain a copy of the License at
8fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville#
9fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville#      http://www.apache.org/licenses/LICENSE-2.0
10fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville#
11fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# Unless required by applicable law or agreed to in writing, software
12fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# distributed under the License is distributed on an "AS IS" BASIS,
13fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# See the License for the specific language governing permissions and
15fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# limitations under the License.
16fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
17fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# This file is used for testing.  The original is at:
18fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville#   http://code.google.com/p/pymox/
19fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
20fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville"""Mox, an object-mocking framework for Python.
21fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
22fbaaef999ba563838ebd00874ed8a1c01fbf286dWink SavilleMox works in the record-replay-verify paradigm.  When you first create
23fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savillea mock object, it is in record mode.  You then programmatically set
24fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savillethe expected behavior of the mock object (what methods are to be
25fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savillecalled on it, with what parameters, what they should return, and in
26fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savillewhat order).
27fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
28fbaaef999ba563838ebd00874ed8a1c01fbf286dWink SavilleOnce you have set up the expected mock behavior, you put it in replay
29fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savillemode.  Now the mock responds to method calls just as you told it to.
30fbaaef999ba563838ebd00874ed8a1c01fbf286dWink SavilleIf an unexpected method (or an expected method with unexpected
31fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleparameters) is called, then an exception will be raised.
32fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
33fbaaef999ba563838ebd00874ed8a1c01fbf286dWink SavilleOnce you are done interacting with the mock, you need to verify that
34fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleall the expected interactions occured.  (Maybe your code exited
35fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleprematurely without calling some cleanup method!)  The verify phase
36fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleensures that every expected method was called; otherwise, an exception
37fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savillewill be raised.
38fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
39fbaaef999ba563838ebd00874ed8a1c01fbf286dWink SavilleSuggested usage / workflow:
40fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
41fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  # Create Mox factory
42fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  my_mox = Mox()
43fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
44fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  # Create a mock data access object
45fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  mock_dao = my_mox.CreateMock(DAOClass)
46fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
47fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  # Set up expected behavior
48fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  mock_dao.RetrievePersonWithIdentifier('1').AndReturn(person)
49fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  mock_dao.DeletePerson(person)
50fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
51fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  # Put mocks in replay mode
52fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  my_mox.ReplayAll()
53fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
54fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  # Inject mock object and run test
55fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  controller.SetDao(mock_dao)
56fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  controller.DeletePersonById('1')
57fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
58fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  # Verify all methods were called as expected
59fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  my_mox.VerifyAll()
60fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville"""
61fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
62fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savillefrom collections import deque
63fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleimport re
64fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleimport types
65fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleimport unittest
66fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
67fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleimport stubout
68fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
69fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass Error(AssertionError):
70fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Base exception for this module."""
71fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
72fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  pass
73fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
74fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
75fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass ExpectedMethodCallsError(Error):
76fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Raised when Verify() is called before all expected methods have been called
77fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
78fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
79fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, expected_methods):
80fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Init exception.
81fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
82fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
83fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # expected_methods: A sequence of MockMethod objects that should have been
84fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      #   called.
85fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      expected_methods: [MockMethod]
86fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
87fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Raises:
88fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      ValueError: if expected_methods contains no methods.
89fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
90fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
91fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if not expected_methods:
92fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      raise ValueError("There must be at least one expected method")
93fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Error.__init__(self)
94fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._expected_methods = expected_methods
95fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
96fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __str__(self):
97fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    calls = "\n".join(["%3d.  %s" % (i, m)
98fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville                       for i, m in enumerate(self._expected_methods)])
99fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return "Verify: Expected methods never called:\n%s" % (calls,)
100fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
101fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
102fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass UnexpectedMethodCallError(Error):
103fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Raised when an unexpected method is called.
104fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
105fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  This can occur if a method is called with incorrect parameters, or out of the
106fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  specified order.
107fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
108fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
109fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, unexpected_method, expected):
110fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Init exception.
111fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
112fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
113fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # unexpected_method: MockMethod that was called but was not at the head of
114fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      #   the expected_method queue.
115fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # expected: MockMethod or UnorderedGroup the method should have
116fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      #   been in.
117fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      unexpected_method: MockMethod
118fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      expected: MockMethod or UnorderedGroup
119fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
120fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
121fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Error.__init__(self)
122fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._unexpected_method = unexpected_method
123fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._expected = expected
124fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
125fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __str__(self):
126fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return "Unexpected method call: %s.  Expecting: %s" % \
127fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      (self._unexpected_method, self._expected)
128fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
129fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
130fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass UnknownMethodCallError(Error):
131fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Raised if an unknown method is requested of the mock object."""
132fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
133fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, unknown_method_name):
134fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Init exception.
135fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
136fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
137fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # unknown_method_name: Method call that is not part of the mocked class's
138fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      #   public interface.
139fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      unknown_method_name: str
140fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
141fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
142fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Error.__init__(self)
143fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._unknown_method_name = unknown_method_name
144fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
145fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __str__(self):
146fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return "Method called is not a member of the object: %s" % \
147fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      self._unknown_method_name
148fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
149fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
150fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass Mox(object):
151fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Mox: a factory for creating mock objects."""
152fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
153fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  # A list of types that should be stubbed out with MockObjects (as
154fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  # opposed to MockAnythings).
155fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  _USE_MOCK_OBJECT = [types.ClassType, types.InstanceType, types.ModuleType,
156fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville                      types.ObjectType, types.TypeType]
157fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
158fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self):
159fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Initialize a new Mox."""
160fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
161fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._mock_objects = []
162fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self.stubs = stubout.StubOutForTesting()
163fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
164fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def CreateMock(self, class_to_mock):
165fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Create a new mock object.
166fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
167fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
168fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # class_to_mock: the class to be mocked
169fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      class_to_mock: class
170fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
171fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
172fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      MockObject that can be used as the class_to_mock would be.
173fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
174fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
175fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    new_mock = MockObject(class_to_mock)
176fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._mock_objects.append(new_mock)
177fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return new_mock
178fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
179fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def CreateMockAnything(self):
180fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Create a mock that will accept any method calls.
181fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
182fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    This does not enforce an interface.
183fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
184fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
185fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    new_mock = MockAnything()
186fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._mock_objects.append(new_mock)
187fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return new_mock
188fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
189fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def ReplayAll(self):
190fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Set all mock objects to replay mode."""
191fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
192fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    for mock_obj in self._mock_objects:
193fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      mock_obj._Replay()
194fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
195fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
196fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def VerifyAll(self):
197fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Call verify on all mock objects created."""
198fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
199fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    for mock_obj in self._mock_objects:
200fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      mock_obj._Verify()
201fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
202fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def ResetAll(self):
203fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Call reset on all mock objects.  This does not unset stubs."""
204fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
205fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    for mock_obj in self._mock_objects:
206fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      mock_obj._Reset()
207fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
208fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def StubOutWithMock(self, obj, attr_name, use_mock_anything=False):
209fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Replace a method, attribute, etc. with a Mock.
210fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
211fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    This will replace a class or module with a MockObject, and everything else
212fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    (method, function, etc) with a MockAnything.  This can be overridden to
213fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    always use a MockAnything by setting use_mock_anything to True.
214fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
215fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
216fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      obj: A Python object (class, module, instance, callable).
217fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      attr_name: str.  The name of the attribute to replace with a mock.
218fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      use_mock_anything: bool. True if a MockAnything should be used regardless
219fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        of the type of attribute.
220fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
221fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
222fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    attr_to_replace = getattr(obj, attr_name)
223fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if type(attr_to_replace) in self._USE_MOCK_OBJECT and not use_mock_anything:
224fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      stub = self.CreateMock(attr_to_replace)
225fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    else:
226fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      stub = self.CreateMockAnything()
227fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
228fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self.stubs.Set(obj, attr_name, stub)
229fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
230fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def UnsetStubs(self):
231fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Restore stubs to their original state."""
232fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
233fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self.stubs.UnsetAll()
234fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
235fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilledef Replay(*args):
236fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Put mocks into Replay mode.
237fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
238fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  Args:
239fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # args is any number of mocks to put into replay mode.
240fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
241fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
242fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  for mock in args:
243fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    mock._Replay()
244fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
245fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
246fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilledef Verify(*args):
247fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Verify mocks.
248fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
249fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  Args:
250fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # args is any number of mocks to be verified.
251fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
252fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
253fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  for mock in args:
254fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    mock._Verify()
255fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
256fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
257fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilledef Reset(*args):
258fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Reset mocks.
259fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
260fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  Args:
261fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # args is any number of mocks to be reset.
262fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
263fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
264fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  for mock in args:
265fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    mock._Reset()
266fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
267fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
268fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass MockAnything:
269fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """A mock that can be used to mock anything.
270fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
271fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  This is helpful for mocking classes that do not provide a public interface.
272fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
273fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
274fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self):
275fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """ """
276fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._Reset()
277fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
278fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __getattr__(self, method_name):
279fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Intercept method calls on this object.
280fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
281fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville     A new MockMethod is returned that is aware of the MockAnything's
282fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville     state (record or replay).  The call will be recorded or replayed
283fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville     by the MockMethod's __call__.
284fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
285fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
286fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # method name: the name of the method being called.
287fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      method_name: str
288fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
289fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
290fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      A new MockMethod aware of MockAnything's state (record or replay).
291fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
292fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
293fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return self._CreateMockMethod(method_name)
294fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
295fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def _CreateMockMethod(self, method_name):
296fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Create a new mock method call and return it.
297fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
298fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
299fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # method name: the name of the method being called.
300fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      method_name: str
301fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
302fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
303fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      A new MockMethod aware of MockAnything's state (record or replay).
304fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
305fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
306fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return MockMethod(method_name, self._expected_calls_queue,
307fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville                      self._replay_mode)
308fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
309fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __nonzero__(self):
310fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Return 1 for nonzero so the mock can be used as a conditional."""
311fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
312fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return 1
313fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
314fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __eq__(self, rhs):
315fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Provide custom logic to compare objects."""
316fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
317fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return (isinstance(rhs, MockAnything) and
318fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville            self._replay_mode == rhs._replay_mode and
319fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville            self._expected_calls_queue == rhs._expected_calls_queue)
320fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
321fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __ne__(self, rhs):
322fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Provide custom logic to compare objects."""
323fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
324fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return not self == rhs
325fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
326fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def _Replay(self):
327fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Start replaying expected method calls."""
328fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
329fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._replay_mode = True
330fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
331fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def _Verify(self):
332fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Verify that all of the expected calls have been made.
333fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
334fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Raises:
335fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      ExpectedMethodCallsError: if there are still more method calls in the
336fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        expected queue.
337fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
338fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
339fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # If the list of expected calls is not empty, raise an exception
340fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if self._expected_calls_queue:
341fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # The last MultipleTimesGroup is not popped from the queue.
342fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      if (len(self._expected_calls_queue) == 1 and
343fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville          isinstance(self._expected_calls_queue[0], MultipleTimesGroup) and
344fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville          self._expected_calls_queue[0].IsSatisfied()):
345fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        pass
346fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      else:
347fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        raise ExpectedMethodCallsError(self._expected_calls_queue)
348fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
349fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def _Reset(self):
350fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Reset the state of this mock to record mode with an empty queue."""
351fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
352fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # Maintain a list of method calls we are expecting
353fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._expected_calls_queue = deque()
354fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
355fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # Make sure we are in setup mode, not replay mode
356fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._replay_mode = False
357fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
358fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
359fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass MockObject(MockAnything, object):
360fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """A mock object that simulates the public/protected interface of a class."""
361fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
362fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, class_to_mock):
363fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Initialize a mock object.
364fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
365fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    This determines the methods and properties of the class and stores them.
366fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
367fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
368fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # class_to_mock: class to be mocked
369fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      class_to_mock: class
370fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
371fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
372fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # This is used to hack around the mixin/inheritance of MockAnything, which
373fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # is not a proper object (it can be anything. :-)
374fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    MockAnything.__dict__['__init__'](self)
375fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
376fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # Get a list of all the public and special methods we should mock.
377fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._known_methods = set()
378fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._known_vars = set()
379fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._class_to_mock = class_to_mock
380fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    for method in dir(class_to_mock):
381fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      if callable(getattr(class_to_mock, method)):
382fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        self._known_methods.add(method)
383fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      else:
384fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        self._known_vars.add(method)
385fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
386fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __getattr__(self, name):
387fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Intercept attribute request on this object.
388fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
389fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    If the attribute is a public class variable, it will be returned and not
390fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    recorded as a call.
391fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
392fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    If the attribute is not a variable, it is handled like a method
393fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    call. The method name is checked against the set of mockable
394fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    methods, and a new MockMethod is returned that is aware of the
395fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    MockObject's state (record or replay).  The call will be recorded
396fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    or replayed by the MockMethod's __call__.
397fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
398fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
399fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # name: the name of the attribute being requested.
400fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      name: str
401fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
402fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
403fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      Either a class variable or a new MockMethod that is aware of the state
404fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      of the mock (record or replay).
405fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
406fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Raises:
407fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      UnknownMethodCallError if the MockObject does not mock the requested
408fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville          method.
409fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
410fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
411fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if name in self._known_vars:
412fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      return getattr(self._class_to_mock, name)
413fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
414fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if name in self._known_methods:
415fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      return self._CreateMockMethod(name)
416fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
417fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    raise UnknownMethodCallError(name)
418fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
419fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __eq__(self, rhs):
420fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Provide custom logic to compare objects."""
421fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
422fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return (isinstance(rhs, MockObject) and
423fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville            self._class_to_mock == rhs._class_to_mock and
424fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville            self._replay_mode == rhs._replay_mode and
425fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville            self._expected_calls_queue == rhs._expected_calls_queue)
426fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
427fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __setitem__(self, key, value):
428fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Provide custom logic for mocking classes that support item assignment.
429fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
430fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
431fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      key: Key to set the value for.
432fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      value: Value to set.
433fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
434fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
435fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      Expected return value in replay mode.  A MockMethod object for the
436fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      __setitem__ method that has already been called if not in replay mode.
437fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
438fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Raises:
439fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      TypeError if the underlying class does not support item assignment.
440fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      UnexpectedMethodCallError if the object does not expect the call to
441fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        __setitem__.
442fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
443fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
444fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    setitem = self._class_to_mock.__dict__.get('__setitem__', None)
445fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
446fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # Verify the class supports item assignment.
447fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if setitem is None:
448fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      raise TypeError('object does not support item assignment')
449fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
450fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # If we are in replay mode then simply call the mock __setitem__ method.
451fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if self._replay_mode:
452fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      return MockMethod('__setitem__', self._expected_calls_queue,
453fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville                        self._replay_mode)(key, value)
454fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
455fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
456fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # Otherwise, create a mock method __setitem__.
457fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return self._CreateMockMethod('__setitem__')(key, value)
458fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
459fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __getitem__(self, key):
460fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Provide custom logic for mocking classes that are subscriptable.
461fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
462fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
463fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      key: Key to return the value for.
464fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
465fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
466fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      Expected return value in replay mode.  A MockMethod object for the
467fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      __getitem__ method that has already been called if not in replay mode.
468fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
469fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Raises:
470fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      TypeError if the underlying class is not subscriptable.
471fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      UnexpectedMethodCallError if the object does not expect the call to
472fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        __setitem__.
473fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
474fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
475fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    getitem = self._class_to_mock.__dict__.get('__getitem__', None)
476fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
477fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # Verify the class supports item assignment.
478fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if getitem is None:
479fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      raise TypeError('unsubscriptable object')
480fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
481fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # If we are in replay mode then simply call the mock __getitem__ method.
482fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if self._replay_mode:
483fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      return MockMethod('__getitem__', self._expected_calls_queue,
484fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville                        self._replay_mode)(key)
485fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
486fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
487fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # Otherwise, create a mock method __getitem__.
488fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return self._CreateMockMethod('__getitem__')(key)
489fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
490fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __call__(self, *params, **named_params):
491fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Provide custom logic for mocking classes that are callable."""
492fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
493fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # Verify the class we are mocking is callable
494fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    callable = self._class_to_mock.__dict__.get('__call__', None)
495fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if callable is None:
496fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      raise TypeError('Not callable')
497fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
498fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # Because the call is happening directly on this object instead of a method,
499fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # the call on the mock method is made right here
500fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    mock_method = self._CreateMockMethod('__call__')
501fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return mock_method(*params, **named_params)
502fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
503fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  @property
504fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __class__(self):
505fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Return the class that is being mocked."""
506fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
507fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return self._class_to_mock
508fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
509fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
510fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass MockMethod(object):
511fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Callable mock method.
512fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
513fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  A MockMethod should act exactly like the method it mocks, accepting parameters
514fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  and returning a value, or throwing an exception (as specified).  When this
515fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  method is called, it can optionally verify whether the called method (name and
516fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  signature) matches the expected method.
517fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
518fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
519fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, method_name, call_queue, replay_mode):
520fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Construct a new mock method.
521fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
522fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
523fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # method_name: the name of the method
524fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # call_queue: deque of calls, verify this call against the head, or add
525fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      #     this call to the queue.
526fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # replay_mode: False if we are recording, True if we are verifying calls
527fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      #     against the call queue.
528fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      method_name: str
529fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      call_queue: list or deque
530fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      replay_mode: bool
531fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
532fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
533fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._name = method_name
534fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._call_queue = call_queue
535fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if not isinstance(call_queue, deque):
536fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      self._call_queue = deque(self._call_queue)
537fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._replay_mode = replay_mode
538fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
539fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._params = None
540fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._named_params = None
541fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._return_value = None
542fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._exception = None
543fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._side_effects = None
544fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
545fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __call__(self, *params, **named_params):
546fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Log parameters and return the specified return value.
547fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
548fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    If the Mock(Anything/Object) associated with this call is in record mode,
549fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    this MockMethod will be pushed onto the expected call queue.  If the mock
550fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    is in replay mode, this will pop a MockMethod off the top of the queue and
551fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    verify this call is equal to the expected call.
552fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
553fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Raises:
554fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      UnexpectedMethodCall if this call is supposed to match an expected method
555fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        call and it does not.
556fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
557fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
558fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._params = params
559fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._named_params = named_params
560fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
561fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if not self._replay_mode:
562fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      self._call_queue.append(self)
563fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      return self
564fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
565fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    expected_method = self._VerifyMethodCall()
566fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
567fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if expected_method._side_effects:
568fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      expected_method._side_effects(*params, **named_params)
569fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
570fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if expected_method._exception:
571fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      raise expected_method._exception
572fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
573fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return expected_method._return_value
574fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
575fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __getattr__(self, name):
576fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Raise an AttributeError with a helpful message."""
577fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
578fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    raise AttributeError('MockMethod has no attribute "%s". '
579fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        'Did you remember to put your mocks in replay mode?' % name)
580fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
581fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def _PopNextMethod(self):
582fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Pop the next method from our call queue."""
583fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    try:
584fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      return self._call_queue.popleft()
585fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    except IndexError:
586fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      raise UnexpectedMethodCallError(self, None)
587fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
588fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def _VerifyMethodCall(self):
589fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Verify the called method is expected.
590fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
591fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    This can be an ordered method, or part of an unordered set.
592fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
593fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
594fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      The expected mock method.
595fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
596fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Raises:
597fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      UnexpectedMethodCall if the method called was not expected.
598fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
599fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
600fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    expected = self._PopNextMethod()
601fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
602fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # Loop here, because we might have a MethodGroup followed by another
603fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # group.
604fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    while isinstance(expected, MethodGroup):
605fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      expected, method = expected.MethodCalled(self)
606fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      if method is not None:
607fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        return method
608fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
609fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # This is a mock method, so just check equality.
610fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if expected != self:
611fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      raise UnexpectedMethodCallError(self, expected)
612fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
613fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return expected
614fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
615fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __str__(self):
616fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    params = ', '.join(
617fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        [repr(p) for p in self._params or []] +
618fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        ['%s=%r' % x for x in sorted((self._named_params or {}).items())])
619fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    desc = "%s(%s) -> %r" % (self._name, params, self._return_value)
620fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return desc
621fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
622fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __eq__(self, rhs):
623fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Test whether this MockMethod is equivalent to another MockMethod.
624fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
625fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
626fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # rhs: the right hand side of the test
627fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      rhs: MockMethod
628fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
629fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
630fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return (isinstance(rhs, MockMethod) and
631fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville            self._name == rhs._name and
632fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville            self._params == rhs._params and
633fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville            self._named_params == rhs._named_params)
634fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
635fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __ne__(self, rhs):
636fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Test whether this MockMethod is not equivalent to another MockMethod.
637fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
638fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
639fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # rhs: the right hand side of the test
640fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      rhs: MockMethod
641fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
642fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
643fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return not self == rhs
644fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
645fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def GetPossibleGroup(self):
646fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Returns a possible group from the end of the call queue or None if no
647fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    other methods are on the stack.
648fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
649fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
650fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # Remove this method from the tail of the queue so we can add it to a group.
651fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    this_method = self._call_queue.pop()
652fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    assert this_method == self
653fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
654fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # Determine if the tail of the queue is a group, or just a regular ordered
655fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # mock method.
656fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    group = None
657fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    try:
658fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      group = self._call_queue[-1]
659fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    except IndexError:
660fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      pass
661fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
662fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return group
663fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
664fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def _CheckAndCreateNewGroup(self, group_name, group_class):
665fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Checks if the last method (a possible group) is an instance of our
666fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    group_class. Adds the current method to this group or creates a new one.
667fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
668fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
669fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
670fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      group_name: the name of the group.
671fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      group_class: the class used to create instance of this new group
672fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
673fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    group = self.GetPossibleGroup()
674fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
675fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # If this is a group, and it is the correct group, add the method.
676fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if isinstance(group, group_class) and group.group_name() == group_name:
677fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      group.AddMethod(self)
678fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      return self
679fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
680fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # Create a new group and add the method.
681fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    new_group = group_class(group_name)
682fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    new_group.AddMethod(self)
683fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._call_queue.append(new_group)
684fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return self
685fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
686fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def InAnyOrder(self, group_name="default"):
687fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Move this method into a group of unordered calls.
688fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
689fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    A group of unordered calls must be defined together, and must be executed
690fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    in full before the next expected method can be called.  There can be
691fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    multiple groups that are expected serially, if they are given
692fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    different group names.  The same group name can be reused if there is a
693fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    standard method call, or a group with a different name, spliced between
694fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    usages.
695fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
696fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
697fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      group_name: the name of the unordered group.
698fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
699fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
700fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      self
701fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
702fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return self._CheckAndCreateNewGroup(group_name, UnorderedGroup)
703fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
704fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def MultipleTimes(self, group_name="default"):
705fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Move this method into group of calls which may be called multiple times.
706fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
707fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    A group of repeating calls must be defined together, and must be executed in
708fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    full before the next expected mehtod can be called.
709fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
710fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
711fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      group_name: the name of the unordered group.
712fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
713fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
714fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      self
715fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
716fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return self._CheckAndCreateNewGroup(group_name, MultipleTimesGroup)
717fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
718fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def AndReturn(self, return_value):
719fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Set the value to return when this method is called.
720fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
721fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
722fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # return_value can be anything.
723fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
724fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
725fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._return_value = return_value
726fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return return_value
727fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
728fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def AndRaise(self, exception):
729fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Set the exception to raise when this method is called.
730fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
731fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
732fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # exception: the exception to raise when this method is called.
733fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      exception: Exception
734fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
735fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
736fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._exception = exception
737fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
738fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def WithSideEffects(self, side_effects):
739fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Set the side effects that are simulated when this method is called.
740fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
741fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
742fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      side_effects: A callable which modifies the parameters or other relevant
743fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        state which a given test case depends on.
744fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
745fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
746fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      Self for chaining with AndReturn and AndRaise.
747fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
748fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._side_effects = side_effects
749fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return self
750fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
751fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass Comparator:
752fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Base class for all Mox comparators.
753fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
754fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  A Comparator can be used as a parameter to a mocked method when the exact
755fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  value is not known.  For example, the code you are testing might build up a
756fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  long SQL string that is passed to your mock DAO. You're only interested that
757fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  the IN clause contains the proper primary keys, so you can set your mock
758fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  up as follows:
759fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
760fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  mock_dao.RunQuery(StrContains('IN (1, 2, 4, 5)')).AndReturn(mock_result)
761fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
762fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  Now whatever query is passed in must contain the string 'IN (1, 2, 4, 5)'.
763fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
764fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  A Comparator may replace one or more parameters, for example:
765fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  # return at most 10 rows
766fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  mock_dao.RunQuery(StrContains('SELECT'), 10)
767fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
768fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  or
769fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
770fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  # Return some non-deterministic number of rows
771fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  mock_dao.RunQuery(StrContains('SELECT'), IsA(int))
772fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
773fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
774fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def equals(self, rhs):
775fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Special equals method that all comparators must implement.
776fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
777fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
778fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      rhs: any python object
779fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
780fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
781fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    raise NotImplementedError, 'method must be implemented by a subclass.'
782fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
783fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __eq__(self, rhs):
784fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return self.equals(rhs)
785fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
786fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __ne__(self, rhs):
787fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return not self.equals(rhs)
788fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
789fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
790fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass IsA(Comparator):
791fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """This class wraps a basic Python type or class.  It is used to verify
792fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  that a parameter is of the given type or class.
793fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
794fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  Example:
795fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  mock_dao.Connect(IsA(DbConnectInfo))
796fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
797fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
798fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, class_name):
799fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Initialize IsA
800fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
801fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
802fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      class_name: basic python type or a class
803fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
804fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
805fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._class_name = class_name
806fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
807fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def equals(self, rhs):
808fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Check to see if the RHS is an instance of class_name.
809fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
810fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
811fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # rhs: the right hand side of the test
812fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      rhs: object
813fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
814fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
815fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      bool
816fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
817fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
818fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    try:
819fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      return isinstance(rhs, self._class_name)
820fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    except TypeError:
821fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # Check raw types if there was a type error.  This is helpful for
822fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # things like cStringIO.StringIO.
823fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      return type(rhs) == type(self._class_name)
824fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
825fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __repr__(self):
826fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return str(self._class_name)
827fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
828fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass IsAlmost(Comparator):
829fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Comparison class used to check whether a parameter is nearly equal
830fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  to a given value.  Generally useful for floating point numbers.
831fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
832fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  Example mock_dao.SetTimeout((IsAlmost(3.9)))
833fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
834fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
835fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, float_value, places=7):
836fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Initialize IsAlmost.
837fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
838fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
839fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      float_value: The value for making the comparison.
840fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      places: The number of decimal places to round to.
841fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
842fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
843fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._float_value = float_value
844fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._places = places
845fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
846fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def equals(self, rhs):
847fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Check to see if RHS is almost equal to float_value
848fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
849fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
850fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      rhs: the value to compare to float_value
851fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
852fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
853fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      bool
854fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
855fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
856fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    try:
857fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      return round(rhs-self._float_value, self._places) == 0
858fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    except TypeError:
859fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # This is probably because either float_value or rhs is not a number.
860fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      return False
861fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
862fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __repr__(self):
863fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return str(self._float_value)
864fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
865fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass StrContains(Comparator):
866fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Comparison class used to check whether a substring exists in a
867fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  string parameter.  This can be useful in mocking a database with SQL
868fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  passed in as a string parameter, for example.
869fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
870fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  Example:
871fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  mock_dao.RunQuery(StrContains('IN (1, 2, 4, 5)')).AndReturn(mock_result)
872fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
873fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
874fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, search_string):
875fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Initialize.
876fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
877fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
878fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # search_string: the string you are searching for
879fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      search_string: str
880fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
881fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
882fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._search_string = search_string
883fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
884fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def equals(self, rhs):
885fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Check to see if the search_string is contained in the rhs string.
886fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
887fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
888fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # rhs: the right hand side of the test
889fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      rhs: object
890fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
891fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
892fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      bool
893fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
894fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
895fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    try:
896fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      return rhs.find(self._search_string) > -1
897fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    except Exception:
898fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      return False
899fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
900fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __repr__(self):
901fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return '<str containing \'%s\'>' % self._search_string
902fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
903fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
904fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass Regex(Comparator):
905fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Checks if a string matches a regular expression.
906fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
907fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  This uses a given regular expression to determine equality.
908fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
909fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
910fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, pattern, flags=0):
911fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Initialize.
912fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
913fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
914fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # pattern is the regular expression to search for
915fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      pattern: str
916fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # flags passed to re.compile function as the second argument
917fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      flags: int
918fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
919fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
920fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self.regex = re.compile(pattern, flags=flags)
921fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
922fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def equals(self, rhs):
923fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Check to see if rhs matches regular expression pattern.
924fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
925fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
926fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      bool
927fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
928fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
929fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return self.regex.search(rhs) is not None
930fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
931fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __repr__(self):
932fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    s = '<regular expression \'%s\'' % self.regex.pattern
933fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if self.regex.flags:
934fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      s += ', flags=%d' % self.regex.flags
935fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    s += '>'
936fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return s
937fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
938fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
939fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass In(Comparator):
940fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Checks whether an item (or key) is in a list (or dict) parameter.
941fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
942fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  Example:
943fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  mock_dao.GetUsersInfo(In('expectedUserName')).AndReturn(mock_result)
944fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
945fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
946fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, key):
947fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Initialize.
948fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
949fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
950fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # key is any thing that could be in a list or a key in a dict
951fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
952fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
953fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._key = key
954fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
955fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def equals(self, rhs):
956fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Check to see whether key is in rhs.
957fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
958fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
959fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      rhs: dict
960fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
961fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
962fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      bool
963fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
964fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
965fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return self._key in rhs
966fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
967fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __repr__(self):
968fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return '<sequence or map containing \'%s\'>' % self._key
969fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
970fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
971fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass ContainsKeyValue(Comparator):
972fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Checks whether a key/value pair is in a dict parameter.
973fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
974fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  Example:
975fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  mock_dao.UpdateUsers(ContainsKeyValue('stevepm', stevepm_user_info))
976fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
977fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
978fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, key, value):
979fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Initialize.
980fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
981fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
982fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # key: a key in a dict
983fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # value: the corresponding value
984fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
985fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
986fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._key = key
987fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._value = value
988fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
989fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def equals(self, rhs):
990fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Check whether the given key/value pair is in the rhs dict.
991fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
992fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
993fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      bool
994fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
995fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
996fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    try:
997fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      return rhs[self._key] == self._value
998fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    except Exception:
999fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      return False
1000fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1001fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __repr__(self):
1002fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return '<map containing the entry \'%s: %s\'>' % (self._key, self._value)
1003fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1004fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1005fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass SameElementsAs(Comparator):
1006fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Checks whether iterables contain the same elements (ignoring order).
1007fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1008fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  Example:
1009fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  mock_dao.ProcessUsers(SameElementsAs('stevepm', 'salomaki'))
1010fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
1011fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1012fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, expected_seq):
1013fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Initialize.
1014fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1015fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
1016fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      expected_seq: a sequence
1017fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
1018fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1019fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._expected_seq = expected_seq
1020fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1021fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def equals(self, actual_seq):
1022fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Check to see whether actual_seq has same elements as expected_seq.
1023fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1024fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
1025fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      actual_seq: sequence
1026fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1027fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
1028fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      bool
1029fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
1030fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1031fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    try:
1032fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      expected = dict([(element, None) for element in self._expected_seq])
1033fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      actual = dict([(element, None) for element in actual_seq])
1034fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    except TypeError:
1035fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # Fall back to slower list-compare if any of the objects are unhashable.
1036fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      expected = list(self._expected_seq)
1037fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      actual = list(actual_seq)
1038fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      expected.sort()
1039fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      actual.sort()
1040fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return expected == actual
1041fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1042fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __repr__(self):
1043fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return '<sequence with same elements as \'%s\'>' % self._expected_seq
1044fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1045fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1046fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass And(Comparator):
1047fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Evaluates one or more Comparators on RHS and returns an AND of the results.
1048fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
1049fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1050fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, *args):
1051fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Initialize.
1052fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1053fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
1054fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      *args: One or more Comparator
1055fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
1056fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1057fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._comparators = args
1058fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1059fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def equals(self, rhs):
1060fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Checks whether all Comparators are equal to rhs.
1061fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1062fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
1063fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # rhs: can be anything
1064fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1065fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
1066fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      bool
1067fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
1068fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1069fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    for comparator in self._comparators:
1070fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      if not comparator.equals(rhs):
1071fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        return False
1072fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1073fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return True
1074fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1075fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __repr__(self):
1076fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return '<AND %s>' % str(self._comparators)
1077fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1078fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1079fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass Or(Comparator):
1080fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Evaluates one or more Comparators on RHS and returns an OR of the results.
1081fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
1082fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1083fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, *args):
1084fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Initialize.
1085fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1086fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
1087fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      *args: One or more Mox comparators
1088fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
1089fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1090fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._comparators = args
1091fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1092fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def equals(self, rhs):
1093fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Checks whether any Comparator is equal to rhs.
1094fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1095fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
1096fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      # rhs: can be anything
1097fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1098fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
1099fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      bool
1100fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
1101fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1102fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    for comparator in self._comparators:
1103fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      if comparator.equals(rhs):
1104fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        return True
1105fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1106fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return False
1107fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1108fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __repr__(self):
1109fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return '<OR %s>' % str(self._comparators)
1110fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1111fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1112fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass Func(Comparator):
1113fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Call a function that should verify the parameter passed in is correct.
1114fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1115fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  You may need the ability to perform more advanced operations on the parameter
1116fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  in order to validate it.  You can use this to have a callable validate any
1117fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  parameter. The callable should return either True or False.
1118fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1119fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1120fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  Example:
1121fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1122fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def myParamValidator(param):
1123fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # Advanced logic here
1124fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return True
1125fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1126fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  mock_dao.DoSomething(Func(myParamValidator), true)
1127fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
1128fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1129fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, func):
1130fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Initialize.
1131fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1132fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
1133fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      func: callable that takes one parameter and returns a bool
1134fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
1135fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1136fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._func = func
1137fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1138fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def equals(self, rhs):
1139fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Test whether rhs passes the function test.
1140fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1141fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    rhs is passed into func.
1142fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1143fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
1144fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      rhs: any python object
1145fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1146fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
1147fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      the result of func(rhs)
1148fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
1149fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1150fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return self._func(rhs)
1151fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1152fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __repr__(self):
1153fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return str(self._func)
1154fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1155fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1156fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass IgnoreArg(Comparator):
1157fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Ignore an argument.
1158fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1159fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  This can be used when we don't care about an argument of a method call.
1160fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1161fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  Example:
1162fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  # Check if CastMagic is called with 3 as first arg and 'disappear' as third.
1163fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  mymock.CastMagic(3, IgnoreArg(), 'disappear')
1164fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
1165fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1166fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def equals(self, unused_rhs):
1167fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Ignores arguments and returns True.
1168fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1169fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
1170fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      unused_rhs: any python object
1171fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1172fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
1173fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      always returns True
1174fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
1175fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1176fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return True
1177fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1178fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __repr__(self):
1179fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return '<IgnoreArg>'
1180fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1181fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1182fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass MethodGroup(object):
1183fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Base class containing common behaviour for MethodGroups."""
1184fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1185fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, group_name):
1186fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._group_name = group_name
1187fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1188fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def group_name(self):
1189fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return self._group_name
1190fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1191fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __str__(self):
1192fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return '<%s "%s">' % (self.__class__.__name__, self._group_name)
1193fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1194fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def AddMethod(self, mock_method):
1195fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    raise NotImplementedError
1196fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1197fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def MethodCalled(self, mock_method):
1198fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    raise NotImplementedError
1199fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1200fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def IsSatisfied(self):
1201fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    raise NotImplementedError
1202fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1203fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass UnorderedGroup(MethodGroup):
1204fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """UnorderedGroup holds a set of method calls that may occur in any order.
1205fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1206fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  This construct is helpful for non-deterministic events, such as iterating
1207fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  over the keys of a dict.
1208fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
1209fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1210fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, group_name):
1211fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    super(UnorderedGroup, self).__init__(group_name)
1212fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._methods = []
1213fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1214fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def AddMethod(self, mock_method):
1215fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Add a method to this group.
1216fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1217fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
1218fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      mock_method: A mock method to be added to this group.
1219fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
1220fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1221fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._methods.append(mock_method)
1222fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1223fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def MethodCalled(self, mock_method):
1224fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Remove a method call from the group.
1225fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1226fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    If the method is not in the set, an UnexpectedMethodCallError will be
1227fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    raised.
1228fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1229fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
1230fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      mock_method: a mock method that should be equal to a method in the group.
1231fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1232fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
1233fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      The mock method from the group
1234fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1235fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Raises:
1236fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      UnexpectedMethodCallError if the mock_method was not in the group.
1237fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
1238fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1239fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # Check to see if this method exists, and if so, remove it from the set
1240fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # and return it.
1241fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    for method in self._methods:
1242fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      if method == mock_method:
1243fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        # Remove the called mock_method instead of the method in the group.
1244fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        # The called method will match any comparators when equality is checked
1245fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        # during removal.  The method in the group could pass a comparator to
1246fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        # another comparator during the equality check.
1247fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        self._methods.remove(mock_method)
1248fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1249fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        # If this group is not empty, put it back at the head of the queue.
1250fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        if not self.IsSatisfied():
1251fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville          mock_method._call_queue.appendleft(self)
1252fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1253fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        return self, method
1254fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1255fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    raise UnexpectedMethodCallError(mock_method, self)
1256fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1257fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def IsSatisfied(self):
1258fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Return True if there are not any methods in this group."""
1259fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1260fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return len(self._methods) == 0
1261fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1262fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1263fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass MultipleTimesGroup(MethodGroup):
1264fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """MultipleTimesGroup holds methods that may be called any number of times.
1265fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1266fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  Note: Each method must be called at least once.
1267fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1268fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  This is helpful, if you don't know or care how many times a method is called.
1269fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
1270fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1271fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(self, group_name):
1272fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    super(MultipleTimesGroup, self).__init__(group_name)
1273fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._methods = set()
1274fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._methods_called = set()
1275fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1276fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def AddMethod(self, mock_method):
1277fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Add a method to this group.
1278fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1279fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
1280fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      mock_method: A mock method to be added to this group.
1281fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
1282fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1283fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self._methods.add(mock_method)
1284fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1285fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def MethodCalled(self, mock_method):
1286fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Remove a method call from the group.
1287fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1288fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    If the method is not in the set, an UnexpectedMethodCallError will be
1289fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    raised.
1290fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1291fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
1292fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      mock_method: a mock method that should be equal to a method in the group.
1293fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1294fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
1295fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      The mock method from the group
1296fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1297fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Raises:
1298fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      UnexpectedMethodCallError if the mock_method was not in the group.
1299fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
1300fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1301fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # Check to see if this method exists, and if so add it to the set of
1302fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # called methods.
1303fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1304fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    for method in self._methods:
1305fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      if method == mock_method:
1306fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        self._methods_called.add(mock_method)
1307fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        # Always put this group back on top of the queue, because we don't know
1308fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        # when we are done.
1309fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        mock_method._call_queue.appendleft(self)
1310fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        return self, method
1311fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1312fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    if self.IsSatisfied():
1313fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      next_method = mock_method._PopNextMethod();
1314fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      return next_method, None
1315fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    else:
1316fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      raise UnexpectedMethodCallError(mock_method, self)
1317fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1318fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def IsSatisfied(self):
1319fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Return True if all methods in this group are called at least once."""
1320fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # NOTE(psycho): We can't use the simple set difference here because we want
1321fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # to match different parameters which are considered the same e.g. IsA(str)
1322fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # and some string. This solution is O(n^2) but n should be small.
1323fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    tmp = self._methods.copy()
1324fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    for called in self._methods_called:
1325fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      for expected in tmp:
1326fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        if called == expected:
1327fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville          tmp.remove(expected)
1328fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville          if not tmp:
1329fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville            return True
1330fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville          break
1331fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return False
1332fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1333fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1334fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass MoxMetaTestBase(type):
1335fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Metaclass to add mox cleanup and verification to every test.
1336fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1337fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  As the mox unit testing class is being constructed (MoxTestBase or a
1338fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  subclass), this metaclass will modify all test functions to call the
1339fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  CleanUpMox method of the test class after they finish. This means that
1340fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  unstubbing and verifying will happen for every test with no additional code,
1341fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  and any failures will result in test failures as opposed to errors.
1342fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
1343fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1344fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def __init__(cls, name, bases, d):
1345fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    type.__init__(cls, name, bases, d)
1346fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1347fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # also get all the attributes from the base classes to account
1348fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    # for a case when test class is not the immediate child of MoxTestBase
1349fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    for base in bases:
1350fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      for attr_name in dir(base):
1351fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        d[attr_name] = getattr(base, attr_name)
1352fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1353fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    for func_name, func in d.items():
1354fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      if func_name.startswith('test') and callable(func):
1355fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        setattr(cls, func_name, MoxMetaTestBase.CleanUpTest(cls, func))
1356fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1357fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  @staticmethod
1358fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def CleanUpTest(cls, func):
1359fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """Adds Mox cleanup code to any MoxTestBase method.
1360fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1361fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Always unsets stubs after a test. Will verify all mocks for tests that
1362fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    otherwise pass.
1363fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1364fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Args:
1365fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      cls: MoxTestBase or subclass; the class whose test method we are altering.
1366fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      func: method; the method of the MoxTestBase test class we wish to alter.
1367fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1368fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    Returns:
1369fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      The modified method.
1370fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    """
1371fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    def new_method(self, *args, **kwargs):
1372fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      mox_obj = getattr(self, 'mox', None)
1373fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      cleanup_mox = False
1374fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      if mox_obj and isinstance(mox_obj, Mox):
1375fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        cleanup_mox = True
1376fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      try:
1377fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        func(self, *args, **kwargs)
1378fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      finally:
1379fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        if cleanup_mox:
1380fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville          mox_obj.UnsetStubs()
1381fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville      if cleanup_mox:
1382fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville        mox_obj.VerifyAll()
1383fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    new_method.__name__ = func.__name__
1384fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    new_method.__doc__ = func.__doc__
1385fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    new_method.__module__ = func.__module__
1386fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    return new_method
1387fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1388fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1389fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass MoxTestBase(unittest.TestCase):
1390fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """Convenience test class to make stubbing easier.
1391fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1392fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  Sets up a "mox" attribute which is an instance of Mox - any mox tests will
1393fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  want this. Also automatically unsets any stubs and verifies that all mock
1394fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  methods have been called at the end of each test, eliminating boilerplate
1395fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  code.
1396fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  """
1397fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1398fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  __metaclass__ = MoxMetaTestBase
1399fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville
1400fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville  def setUp(self):
1401fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville    self.mox = Mox()
1402