Mock Server for Django Integration Tests

mock_server.py

import httplib
import multiprocessing
import os
import re
from collections import defaultdict
from wsgiref.simple_server import make_server

from django.core.handlers.wsgi import WSGIHandler
from django.http import HttpResponse


class MockApplication(WSGIHandler):
    """A Mock Application handler.

    This loosly follows the flask application api, but uses Django's Request
    and Response objects as well as the internal wsgi handler. A little hacky
    but hey this is just a Mock Server and it is not intended for production.
    
    Usage:
    
        my_fake_service = MockApplication()
        
        @my_fake_service.route('/v1/faker')
        def fakey_fake():
            return json.dumps({'a': 'fake response'})
        
        
        class MyTest(TestCase):
        
             @classmethod
             def setUpClass(cls):
                 cls.app = my_fake_service
                 cls.process = multiprocessing.Process(
                     target=self.app.run,
                     kwargs={'port': 8080, 'threaded': True}
                 )
                 cls.process.daemon = True
                 cls.process.start()
                 time.sleep(0.5)  # let it start up, might not be needed :shrug:
             
             @classmethod
             def tearDownClass(cls):
                 os.kill(cls.process.pid, signal.SIGKILL)
             
             def test_something_that_uses_mock_services(self):
                 resp = request.post('http://localhost:8080/v1/faker')
                 self.assertEquals(resp.json(), {'a': 'fake response'})
    """

    def __init__(self, queue=None):
        self.queue = queue or multiprocessing.Queue()
        # The set of URLs for this application
        # { 'POST': { '/a/path/to/match': view_function_to_call } }
        self.urls = defaultdict(dict)
        super(MockApplication, self).__init__()

    def route(self, path, methods=['GET']):
        """
        A simplified route decorator to expose a route like flask:

            app = MockApplication()

            @app.route('/hello')
            def hello_view():
                return 'hello world'

        You can specify the path as a regex to pass arguments to the view:

            @app.route('/v1/people/(?P<id>\w+)', methods=['GET', 'POST'])
            def person_handler(id):
                # GET /v1/people/12345 -> person_handler(id=12345)
                return json.dumps({'person': {'id': id}})
        """
        def _decorator(function):
            for method in methods:
                self.urls[method][path] = function
        return _decorator

    def run(self, port, **kwargs):
        """Run the application. Called by `MockServer.start()`"""
        httpd = make_server('', int(port), self)
        httpd.serve_forever()

    def _get_headers(self, request):
        """
        Turn Django request headers into common headers.

        Django stores headers in a `META` dict like `HTTP_X_AUTH_TOKEN` but
        the rest of the world uses headers like `X-Auth-Token`. This just
        converts those to ones we are expecting in our tests.
        """
        headers = {}
        for key, value in request.META.items():
            if key.startswith('HTTP'):
                header = key.replace('HTTP_', '').replace('_', '-').title()
                headers[header] = value
        return headers

    def get_response(self, request):
        """
        This method is called by the base handler and normally uses the
        Django urlconf to dispatch the request to the correct view. We
        have hijacked routing with our internal 'route' decorator.
        """

        # Flask normally sets the content type with 'jsonify' response helper
        # but nearly all of our mock servers return json. We just set this by
        # default. If it needs to be changed simply alter the view response to
        # use the `tuple` response noted below in 'Handle Response'.
        headers = {"Content-Type": "application/json"}

        # These are optional values that could be set by the view.
        extra_headers = {}
        code = 200

        # The default response in case no routes are matched from the request.
        response = ('{"error":"you have failed"}', 404)

        for route, view in self.urls[request.method].items():
            match = re.match(route, request.path)
            if match is None:
                continue

            try:
                # Some views need the request, flask 'cheats' by using
                # threadlocals, here we just try passing the request object
                # we will get a TypeError if the function does not accept
                # a request argument.
                #
                # We also pass **match.groupdict() to catch any regular
                # expression groups that were defined with the route decorator.
                response = view(request, **match.groupdict())
            except TypeError:
                response = view(**match.groupdict())

        # Handle Response
        # ---------------
        # Flask automatically creates a response object from the views. This
        # makes it really simple to define view methods. By default if your
        # view returns a string that is considered the content of the response.
        # and the status code is set to 200 by default. If you need to modify
        # the status code or add headers to the response Flask allows you to
        # return a tuple. The tuple can either be (content, status_code) or it
        # can be (content, status_code, headers)
        #
        # Most of our mocks just return the output of json.dumps() but a few
        # take advantage of the tuple response tricks. First we check for these
        # special cases first and adjust the status code and headers accordingly.
        if isinstance(response, tuple):
            if len(response) == 3:
                content, code, extra_headers = response
            elif len(response) == 2:
                content, code = response
        else:
            content = response

        resp = HttpResponse(content, status=code)

        # Apply all the headers to the response. Django forces you to set them
        # via the __setitem__ method aka dict[key] = value.
        headers.update(extra_headers)
        for header, value in headers.items():
            resp[header] = value

        # Record the request to allow tests to check it was called.
        self.queue.put({
            "url": request.path,
            "method": request.method,
            "body": content,
            "headers": self._get_headers(request)
        })

        return resp