10f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles)# Copyright 2013 The Chromium Authors. All rights reserved. 20f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles)# Use of this source code is governed by a BSD-style license that can be 30f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles)# found in the LICENSE file. 40f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) 50f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles)class MockFunction(object): 60f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) '''Decorates a function to record the number of times it's called, and 70f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) use that to make test assertions. 80f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) 90f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) Use like: 100f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) 110f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) @MockFunction 120f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) def my_function(): pass 130f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) my_function() 140f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) my_function() 150f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) self.assertTrue(*my_function.CheckAndReset(2)) 160f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) 170f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) or 180f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) 190f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) my_constructor = MockFunction(HTMLParser) 200f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) my_constructor() 210f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) self.assertTrue(*my_constructor.CheckAndReset(1)) 220f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) 230f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) and so on. 240f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) ''' 250f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) 260f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) def __init__(self, fn): 270f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) self._fn = fn 280f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) self._call_count = 0 290f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) 300f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) def __call__(self, *args, **optargs): 310f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) self._call_count += 1 320f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) return self._fn(*args, **optargs) 330f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) 340f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) def CheckAndReset(self, expected_call_count): 350f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) actual_call_count = self._call_count 360f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) self._call_count = 0 370f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) if expected_call_count == actual_call_count: 380f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) return True, '' 390f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) return (False, '%s: expected %s call(s), got %s' % 400f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles) (self._fn.__name__, expected_call_count, actual_call_count)) 41