Create fixtures for urllib.request.urlopen()

gistfile1.py

from io import StringIO
import urllib.request as urllib2
import imp

class MockHTTP(object):

    def __init__(self):
        class MyHTTPHandler(urllib2.HTTPHandler):
            def http_open(handler, req):
                return self.mock_response(req)
        self.expect_url = {}
        my_opener = urllib2.build_opener(MyHTTPHandler)
        urllib2.install_opener(my_opener)

    def generate_response(self, req, return_data, return_code, return_msg):
        resp = urllib2.addinfourl(StringIO(return_data),
                                    'mock header',
                                    req.get_full_url(), code=return_code)
        resp.msg = return_msg
        return resp

    def register(self, url, return_data, return_code=200, return_msg='OK'):
        self.expect_url[url] = lambda req:self.generate_response(req, return_data,
                                                        return_code, return_msg)

    def mock_response(self, req):
        return self.expect_url[req.get_full_url()](req)

    def stop(self):
        imp.reload(urllib2)

from nose.tools import *
from urllib.error import HTTPError

class TestMockHTTP(object):

    def setup(self):
        self.mh = MockHTTP()

    def teardown(self):
        self.mh = MockHTTP()

    """ urlopen should return 'hello' """
    def test_urlopen(self):
        self.mh.register('http://mock_url.example', 'hello')
        eq_(urllib2.urlopen('http://mock_url.example').read(), 'hello')

    """ raise HTTPError if return code is not 200 """
    def test_return_code(self):
        self.mh.register('http://mock_url.example', 'hello', 
                         return_code=404, return_msg='Not Found')
        try:
            eq_(urllib2.urlopen('http://mock_url.example').read(), 'hello')
        except HTTPError as e:
            eq_(404, e.getcode(), msg='HTTP error must be Not Found')

    """ urlopen should raise HTTPError after mock http is stopped """
    def test_stop_mock(self):
        self.mh.stop()
        try:
            urllib2.urlopen('http://hoge')
        except HTTPError as e:
            eq_(504, e.getcode(), msg='HTTP error must be Gateway Time-out')