1# Copyright 2013 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5class MockFunction(object):
6  '''Decorates a function to record the number of times it's called, and
7  use that to make test assertions.
8
9  Use like:
10
11  @MockFunction
12  def my_function(): pass
13  my_function()
14  my_function()
15  self.assertTrue(*my_function.CheckAndReset(2))
16
17  or
18
19  my_constructor = MockFunction(HTMLParser)
20  my_constructor()
21  self.assertTrue(*my_constructor.CheckAndReset(1))
22
23  and so on.
24  '''
25
26  def __init__(self, fn):
27    self._fn = fn
28    self._call_count = 0
29
30  def __call__(self, *args, **optargs):
31    self._call_count += 1
32    return self._fn(*args, **optargs)
33
34  def CheckAndReset(self, expected_call_count):
35    actual_call_count = self._call_count
36    self._call_count = 0
37    if expected_call_count == actual_call_count:
38      return True, ''
39    return (False, '%s: expected %s call(s), got %s' %
40                   (self._fn.__name__, expected_call_count, actual_call_count))
41