monkey_example.py

gistfile1.py

#coding: utf-8
import socket

class FooClient(object):
    def __init__(self, host, port):
        self._sock = socket.socket()
        self._sock.connect((host, port))

    def say(self):
        self._sock.sendall('Foo!\n')


# test library

class Monkey(object):
    def __init__(self):
        self._backup = []

    def patch(self, obj, name, attr):
        ins = getattr(obj, name)
        self._backup.append((obj, name, ins))
        setattr(obj, name, attr)

    def restore(self):
        for obj, name, ins in reversed(self._backup):
            try:
                setattr(obj, name, ins)
            except Exception:
                pass
        del self._backup[:]

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        self.restore()

# test code

def test_foo():
    class dummy_socket(object):
        def __init__(self):
            self._written = []
            self._peer = None

        def connect(self, peer):
            self._peer = peer

        def sendall(self, msg):
            self._written.append(msg)

    with Monkey() as monkey: # with ブロックを抜けると monkey.__exit__() が呼ばれる.
        monkey.patch(socket, 'socket', dummy_socket)
        foo = FooClient('dummy', 1234)
        foo.say()
        foo.say()
        assert foo._sock._peer == ('dummy', 1234)
        assert foo._sock._written == ['Foo!\n']*2

test_foo()
print socket.socket # もとに戻ってる