httmock.py

from functools import wraps
import datetime
from requests import cookies
import json
import re
import requests
from requests import structures, utils
import sys
try:
    import urlparse
except ImportError:
    import urllib.parse as urlparse

if sys.version_info >= (3, 0, 0):
    from io import BytesIO
else:
    from StringIO import StringIO as BytesIO


binary_type = bytes
if sys.version_info >= (3, 0, 0):
    text_type = str
else:
    text_type = unicode  # noqa


class Headers(object):
    def __init__(self, res):
        self.headers = res.headers

    def get_all(self, name, failobj=None):
        return self.getheaders(name)

    def getheaders(self, name):
        return [self.headers.get(name)]


def response(status_code=200, content='', headers=None, reason=None, elapsed=0,
             request=None, stream=False, http_vsn=11):
    res = requests.Response()
    res.status_code = status_code
    if isinstance(content, (dict, list)):
        content = json.dumps(content).encode('utf-8')
    if isinstance(content, text_type):
        content = content.encode('utf-8')
    res._content = content
    res._content_consumed = content
    res.headers = structures.CaseInsensitiveDict(headers or {})
    res.encoding = utils.get_encoding_from_headers(res.headers)
    res.reason = reason
    res.elapsed = datetime.timedelta(elapsed)
    res.request = request
    if hasattr(request, 'url'):
        res.url = request.url
        if isinstance(request.url, bytes):
            res.url = request.url.decode('utf-8')
    if 'set-cookie' in res.headers:
        res.cookies.extract_cookies(cookies.MockResponse(Headers(res)),
                                    cookies.MockRequest(request))
    if stream:
        res.raw = BytesIO(content)
    else:
        res.raw = BytesIO(b'')
    res.raw.version = http_vsn

    # normally this closes the underlying connection,
    #  but we have nothing to free.
    res.close = lambda *args, **kwargs: None

    return res


def all_requests(func):
    @wraps(func)
    def inner(*args, **kwargs):
        return func(*args, **kwargs)
    return inner


def urlmatch(scheme=None, netloc=None, path=None, method=None, query=None):
    def decorator(func):
        @wraps(func)
        def inner(self_or_url, url_or_request, *args, **kwargs):
            if isinstance(self_or_url, urlparse.SplitResult):
                url = self_or_url
                request = url_or_request
            else:
                url = url_or_request
                request = args[0]
            if scheme is not None and scheme != url.scheme:
                return
            if netloc is not None and not re.match(netloc, url.netloc):
                return
            if path is not None and not re.match(path, url.path):
                return
            if query is not None and not re.match(query, url.query):
                return
            if method is not None and method.upper() != request.method:
                return
            return func(self_or_url, url_or_request, *args, **kwargs)
        return inner
    return decorator


def handler_init_call(handler):
    setattr(handler, 'call', {
        'count': 0,
        'called': False,
        'requests': []
    })


def handler_clean_call(handler):
    if hasattr(handler, 'call'):
        handler.call.update({
            'count': 0,
            'called': False,
            'requests': []
        })


def handler_called(handler, *args, **kwargs):
    try:
        return handler(*args, **kwargs)
    finally:
        handler.call['count'] += 1
        handler.call['called'] = True
        handler.call['requests'].append(args[1])

def remember_called(func):
    handler_init_call(func)

    @wraps(func)
    def inner(*args, **kwargs):
        return handler_called(func, *args, **kwargs)
    return inner


def first_of(handlers, *args, **kwargs):
    for handler in handlers:
        res = handler(*args, **kwargs)
        if res is not None:
            return res


class HTTMock(object):
    """
    Acts as a context manager to allow mocking
    """
    STATUS_CODE = 200

    def __init__(self, *handlers):
        self.handlers = handlers

    def __enter__(self):
        self._real_session_send = requests.Session.send
        self._real_session_prepare_request = requests.Session.prepare_request

        for handler in self.handlers:
            handler_clean_call(handler)

        def _fake_send(session, request, **kwargs):
            response = self.intercept(request, **kwargs)

            if isinstance(response, requests.Response):
                # this is pasted from requests to handle redirects properly:
                kwargs.setdefault('stream', session.stream)
                kwargs.setdefault('verify', session.verify)
                kwargs.setdefault('cert', session.cert)
                kwargs.setdefault('proxies', session.proxies)

                allow_redirects = kwargs.pop('allow_redirects', True)
                stream = kwargs.get('stream')
                timeout = kwargs.get('timeout')
                verify = kwargs.get('verify')
                cert = kwargs.get('cert')
                proxies = kwargs.get('proxies')

                gen = session.resolve_redirects(
                    response,
                    request,
                    stream=stream,
                    timeout=timeout,
                    verify=verify,
                    cert=cert,
                    proxies=proxies)

                history = [resp for resp in gen] if allow_redirects else []

                if history:
                    history.insert(0, response)
                    response = history.pop()
                    response.history = tuple(history)

                session.cookies.update(response.cookies)

                return response

            return self._real_session_send(session, request, **kwargs)

        def _fake_prepare_request(session, request):
            """
            Fake this method so the `PreparedRequest` objects contains
            an attribute `original` of the original request.
            """
            prep = self._real_session_prepare_request(session, request)
            prep.original = request
            return prep

        requests.Session.send = _fake_send
        requests.Session.prepare_request = _fake_prepare_request

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        requests.Session.send = self._real_session_send
        requests.Session.prepare_request = self._real_session_prepare_request

    def intercept(self, request, **kwargs):
        url = urlparse.urlsplit(request.url)
        res = first_of(self.handlers, url, request)

        if isinstance(res, requests.Response):
            return res
        elif isinstance(res, dict):
            return response(res.get('status_code'),
                            res.get('content'),
                            res.get('headers'),
                            res.get('reason'),
                            res.get('elapsed', 0),
                            request,
                            stream=kwargs.get('stream', False),
                            http_vsn=res.get('http_vsn', 11))
        elif isinstance(res, (text_type, binary_type)):
            return response(content=res, stream=kwargs.get('stream', False))
        elif res is None:
            return None
        else:
            raise TypeError(
                "Dont know how to handle response of type {0}".format(type(res)))


def with_httmock(*handlers):
    mock = HTTMock(*handlers)

    def decorator(func):
        @wraps(func)
        def inner(*args, **kwargs):
            with mock:
                return func(*args, **kwargs)
        return inner
    return decorator