urllib.request の fixture

HTTPretty が python3.3 の urllib.request.urlopen() でうまく動かなかったので探してきました。で参考にさせて頂いたソースがあったのですが、どこか忘れてしまった・・・手順は以下のとおり。urllib2 を使えば python2 でもたぶん大丈夫。

仕込み

from io import StringIO
import urllib.request as urllib2

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

expect_url = {}

def register(url, return_data):
    expect_url[url] = lambda req:generate_response(req, return_data)

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

class MyHTTPHandler(urllib2.HTTPHandler):
    def http_open(self, req):
        return mock_response(req)

これを install_opener します。あとは、expect_url にハックする URL と戻り値を入れればいいです。

テスト本体

from nose.tools import *

class TestUrlopen(object):
    def __init__(self):
        my_opener = urllib2.build_opener(MyHTTPHandler)
        urllib2.install_opener(my_opener)

    def test_urlopen(self):
        url = 'http://mock_server'
        register(url, "hello world")

        ret = urllib2.urlopen(url).read()
        eq_(ret, "hello world")

expect_url を global に置いて動的にレスポンスをつけるのは私のアイデアですが、他はまんまコードを頂いてます。どこだったかな・・・

ここではやってませんが、同じURLで違うデータを返す（副作用を入れる)場合は register() で push して、mock_response() で pop すれば良いかと。

なんとなくクラスも作ってみました。https://gist.github.com/norobust/8254063