モンキーパッチを管理する

テストとかで単純にパッチを管理する. with文を使って呼び出すことで確実にパッチを除去することができる. スレッドセーフでないので注意.

patch.py

# -*- coding: utf-8

from unittest import TestCase

class PatchManager(object):
    def __init__(self):
        self._storage = {}

    def attach(self, target, name, patch):
        key = '%s.%s'%(target.__name__, name)
        if self._storage.has_key(key):
            raise ValueError('Patch for %s already exists.' % key)
        self._storage[key] = {
            'target': target,
            'name': name,
            'original': getattr(target, name)
            }
        setattr(target, name, patch)
        return

    def detach(self, target_name):
        original = self._storage.get(target_name)
        setattr(original['target'], original['name'], original['original'])
        return

    def detach_all(self):
        for key in self._storage.keys():
            self.detach(key)
        return

    def __enter__(self):
        return self

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

class TestPatchManager(TestCase):

    def test_1(self):
        """Simple patching."""

        import datetime

        def fake():
            return 'fake'

        origin = datetime.datetime

        mgr = PatchManager()
        mgr.attach(datetime, 'datetime', fake)

        self.assertNotEqual(datetime.datetime, origin)
        self.assertEqual(datetime.datetime, fake)

        mgr.detach('datetime.datetime')
        self.assertEqual(datetime.datetime, origin)

        return

    def test_2(self):
        """Call detach all."""

        import datetime

        def fake():
            return 'fake'

        origin = datetime.datetime

        mgr = PatchManager()
        mgr.attach(datetime, 'datetime', fake)

        self.assertNotEqual(datetime.datetime, origin)
        self.assertEqual(datetime.datetime, fake)

        mgr.detach_all()
        self.assertEqual(datetime.datetime, origin)

        return

    def test_3(self):
        """Use with with context."""

        import datetime

        def fake():
            return 'fake'

        origin = datetime.datetime

        with PatchManager() as mgr:
            mgr.attach(datetime, 'datetime', fake)

            self.assertNotEqual(datetime.datetime, origin)
            self.assertEqual(datetime.datetime, fake)

        self.assertEqual(datetime.datetime, origin)

        return

    def test_4(self):
        """If exceptions is occured in with context."""

        import datetime

        def fake():
            return 'fake'

        origin = datetime.datetime

        with self.assertRaises(AssertionError):
            with PatchManager() as mgr:
                mgr.attach(datetime, 'datetime', fake)

                self.assertNotEqual(datetime.datetime, origin)
                self.assertEqual(datetime.datetime, fake)

                assert False

        self.assertEqual(datetime.datetime, origin)

        return

    def test_5(self):
        """Patched to the same obuject twice."""

        import datetime

        def fake():
            return 'fake'

        origin = datetime.datetime

        with self.assertRaises(ValueError):
            with PatchManager() as mgr:
                mgr.attach(datetime, 'datetime', fake)
                mgr.attach(datetime, 'datetime', fake)

                assert False

        self.assertEqual(datetime.datetime, origin)

        return