Google AppEngine URLFetch in Unit Tests

gistfile1.py

from google.appengine.api import apiproxy_stub
from google.appengine.api import apiproxy_stub_map

class FetchServiceMock(apiproxy_stub.APIProxyStub):
    def __init__(self, service_name='urlfetch'):
        super(FetchServiceMock, self).__init__(service_name)

    def set_return_values(self, **kwargs):
        self.return_values = kwargs

    def _Dynamic_Fetch(self, request, response):
        rv = self.return_values
        response.set_content(rv.get('content', ''))
        response.set_statuscode(rv.get('status_code', 500))
        for header_key, header_value in rv.get('headers', {}):
            new_header = response.add_header() # prototype for a header
            new_header.set_key(header_key)
            new_header.set_value(header_value)
        response.set_finalurl(rv.get('final_url', request.url))
        response.set_contentwastruncated(rv.get('content_was_truncated', False))

        # allow to query the object after it is used
        self.request = request
        self.response = response

class TestSomeApi(unittest.TestCase):
    def setUp(self):
        # initialize apiproxy
        self.fetch_mock = FetchServiceMock()
        apiproxy_stub_map.apiproxy.RegisterStub('urlfetch', self.fetch_mock)

    def test_my_method(self):
        self.fetch_mock.set_return_values(
            status_code=200,
            content="success"
        )
        # call some method of the api
        stuff = SomeApi().get_some_important_stuff(goody_company)
        self.assertEquals("http://good.com/stuff", self.fetch_mock.request.url())
        self.assertEquals("success", stuff.data)

I forked it and turned it into a context manager.