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 5 class 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
