Count / verify the number of method calls.

It relies on the snippet I mentioned earlier (http://qiita.com/items/1321). It counts / validates method calls in tests and so on.

from unittest import TestCase


class CallCounter(object):
    """Check a numbe of time a callable is called."""

    def __new__(cls, *args, **kargs):
        inst = super(CallCounter, cls).__new__(cls, *args, **kargs)
        inst._patch_mgr = PatchManager()
        inst._targets = {}
        return inst

    class FailedVerification(Exception):
        pass

    class CallCountManager(object):
        def __init__ (self, target, name, expected):
            self.__target = target
            self.__name = name
            self.__original = getattr(target, name)
            self.__count = 0
            self.__expected = expected

        def __call__(self, *args, **kargs):
            self.__count += 1
            return self.__original(*args, **kargs)

        def __verify(self):
            if self.__count == self.__expected:
                return

            msg = '%s.%s should be called %d times but %d' % \
                  (self.__target.__name__, self.__name, self.__expected,
                   self.__count,)
            raise CallCounter.FailedVerification(msg)

    def ensure(self, target, name, count):
        counter = self.CallCountManager(target, name, count)
        key = '%s.%s'%(target.__name__, name)
        self._targets[key] = counter
        self._patch_mgr.attach(target, name, counter)
        return

    def verify(self):
        for key, counter in self._targets.iteritems():
            counter._CallCountManager__verify()
        return

    def end(self):
        self._patch_mgr.detach_all()
        return

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end()
        if exc_type:
            return False
        return True


class TestCallCounter(TestCase):
    """Test for test.CallCounter."""

    def test_1(self):
        """Call it simply."""

        c = CallCounter()

        import datetime
        original = datetime.datetime
        c.ensure(datetime, 'datetime', 1)
        self.assert_(isinstance(datetime.datetime,
                                CallCounter.CallCountManager))
        self.assertNotEqual(datetime.datetime, original)

        # Call it first.
        datetime.datetime(year=2011, month=3, day=4)
        c.verify()

        # Call it second.
        datetime.datetime(year=2011, month=3, day=4)
        with self.assertRaises(CallCounter.FailedVerification):
            c.verify()

        # Verify a clearance.
        c.end()
        self.assertEqual(datetime.datetime, original)

        return

    def test_2(self):
        """Call it simply."""

        c = CallCounter()

        import datetime
        original = datetime.datetime
        c.ensure(datetime, 'datetime', 0)
        self.assertNotEqual(datetime.datetime, original)

        # Call it and verify.
        datetime.datetime(year=2011, month=3, day=4)
        with self.assertRaises(CallCounter.FailedVerification):
            c.verify()

        c.end()
        self.assertEqual(datetime.datetime, original)

        return

    def test_3(self):
        """Use it with `with statement`."""

        import datetime
        original = datetime.datetime
        with self.assertRaises(CallCounter.FailedVerification):
            with CallCounter() as c:
                c.ensure(datetime, 'datetime', 0)
                self.assertNotEqual(datetime.datetime, original)

                # Call it and verify.
                datetime.datetime(year=2011, month=3, day=4)
                c.verify()

        self.assertEqual(datetime.datetime, original)

        return

    def test_4(self):
        """Exception is happend in with statement."""

        import datetime
        original = datetime.datetime
        with self.assertRaises(AssertionError):
            with CallCounter() as c:
                c.ensure(datetime, 'datetime', 0)
                self.assertNotEqual(datetime.datetime, original)

                # Call it and verify.
                datetime.datetime(year=2011, month=3, day=4)
                raise AssertionError()

        self.assertEqual(datetime.datetime, original)

        return