| # Copyright 2013 The Chromium Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| class MockFunction(object): |
| '''Decorates a function to record the number of times it's called, and |
| use that to make test assertions. |
| |
| Use like: |
| |
| @MockFunction |
| def my_function(): pass |
| my_function() |
| my_function() |
| self.assertTrue(*my_function.CheckAndReset(2)) |
| |
| or |
| |
| my_constructor = MockFunction(HTMLParser) |
| my_constructor() |
| self.assertTrue(*my_constructor.CheckAndReset(1)) |
| |
| and so on. |
| ''' |
| |
| def __init__(self, fn): |
| self._fn = fn |
| self._call_count = 0 |
| |
| def __call__(self, *args, **optargs): |
| self._call_count += 1 |
| return self._fn(*args, **optargs) |
| |
| def CheckAndReset(self, expected_call_count): |
| actual_call_count = self._call_count |
| self._call_count = 0 |
| if expected_call_count == actual_call_count: |
| return True, '' |
| return (False, '%s: expected %s call(s), got %s' % |
| (self._fn.__name__, expected_call_count, actual_call_count)) |