AsyncHTTPStubClient.py

#! /usr/bin/env python
# coding:utf-8
# author: fengchang08@gmail.com

import functools
from io import StringIO
from tornado.httpclient import HTTPRequest, HTTPResponse


class AsyncHTTPStubClient(object):

    def __init__(self):
        self.repsonses = []

    def add_response(self, body, code=200):
        """
        添加一个将要被返回的 HTTPResponse
        :param body: HTTPResponse 的 body, 应该是 unicode
        :param code: HTTPResponse 的状态码, 应该是 int
        """
        response_partial = functools.partial(HTTPResponse, code=code, buffer=StringIO(body))
        self.repsonses.append(response_partial)

    def fetch(self, request, callback):
        """
        返回一个 HTTP Response, 如果没有就抛异常
        :param request: HTTPRequest 或者 URL
        :param callback: 回调函数
        """
        if not isinstance(request, HTTPRequest):
            request = HTTPRequest(url=request)
        if len(self.repsonses) > 0:
            response_partial = self.repsonses.pop(0)
            resp = response_partial(request)
            callback(resp)
        else:
            raise Exception("No more response to return.")