http_server.py

import threading
import json

from collections import defaultdict
from contextlib import contextmanager
from http import HTTPStatus
from http.cookies import SimpleCookie
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs

import pytest


class TestHandler(BaseHTTPRequestHandler):
    handlers = defaultdict(dict)

    @classmethod
    def handler(cls, method, path):
        def inner(func):
            cls.handlers[method][path] = func
            return func
        return inner

    def do_generic(self):
        parse_result = urlparse(self.path)
        func = self.handlers[self.command].get(parse_result.path)
        if func is None:
            return self.send_error(HTTPStatus.NOT_FOUND)

        return func(self)

    do_GET = do_generic
    do_POST = do_generic


@TestHandler.handler('GET', '/headers')
def get_headers(handler):
    handler.send_response(200)
    for key, value in handler.headers.items():
        handler.send_header(key, value)
    handler.send_header('Content-Length', 0)
    handler.end_headers()


def _send_chunked_response(handler, content_type, chunks):
    """Send a chunked HTTP response with the given content type and body chunks."""
    handler.send_response(200)
    if content_type is not None:
        handler.send_header('Content-Type', content_type)
    handler.send_header('Transfer-Encoding', 'chunked')
    handler.end_headers()

    for body in chunks:
        handler.wfile.write(f'{len(body):X}\r\n{body}\r\n'.encode('utf-8'))

    handler.wfile.write('0\r\n\r\n'.encode('utf-8'))


@TestHandler.handler('GET', '/drip')
def chunked_drip(handler):
    _send_chunked_response(
        handler,
        content_type=handler.headers.get('Accept'),
        chunks=['test\n' for _ in range(3)],
    )


@TestHandler.handler('GET', '/ndjson')
def ndjson_stream(handler):
    _send_chunked_response(
        handler,
        content_type='application/x-ndjson',
        chunks=[
            json.dumps({"index": i, "message": "test"}) + '\n'
            for i in range(3)
        ],
    )


@TestHandler.handler('GET', '/sse-json')
def sse_json_stream(handler):
    _send_chunked_response(
        handler,
        content_type='text/event-stream',
        chunks=[
            f'data: {json.dumps({"index": i, "message": "test"})}\n\n'
            for i in range(3)
        ],
    )


@TestHandler.handler('GET', '/stream/encoding/random')
def random_encoding(handler):
    from tests.fixtures import ASCII_FILE_CONTENT, FILE_CONTENT as UNICODE_FILE_CONTENT

    handler.send_response(200)
    handler.send_header('Transfer-Encoding', 'chunked')
    handler.end_headers()

    for body in [
        ASCII_FILE_CONTENT,
        ASCII_FILE_CONTENT,
        UNICODE_FILE_CONTENT,
        UNICODE_FILE_CONTENT,
        UNICODE_FILE_CONTENT,
    ]:
        body += "\n"
        handler.wfile.write(f'{len(body.encode()):X}\r\n{body}\r\n'.encode())

    handler.wfile.write('0\r\n\r\n'.encode('utf-8'))


@TestHandler.handler('POST', '/status/msg')
def status_custom_msg(handler):
    content_len = int(handler.headers.get('content-length', 0))
    post_body = handler.rfile.read(content_len).decode()

    handler.send_response(200, post_body)
    handler.end_headers()


@TestHandler.handler('GET', '/cookies')
def get_cookies(handler):
    cookies = {
        'cookies': {
            key: cookie.value
            for key, cookie in SimpleCookie(handler.headers.get('Cookie')).items()
        }
    }
    payload = json.dumps(cookies)

    handler.send_response(200)
    handler.send_header('Content-Length', len(payload))
    handler.send_header('Content-Type', 'application/json')
    handler.end_headers()
    handler.wfile.write(payload.encode('utf-8'))


@TestHandler.handler('GET', '/cookies/set')
def set_cookies(handler):
    options = parse_qs(urlparse(handler.path).query)

    handler.send_response(200)
    for cookie, [value] in options.items():
        handler.send_header('Set-Cookie', f'{cookie}={value}')

    handler.end_headers()


@TestHandler.handler('GET', '/cookies/set-and-redirect')
def set_cookie_and_redirect(handler):
    handler.send_response(302)

    redirect_to = handler.headers.get('X-Redirect-To', '/headers')
    handler.send_header('Location', redirect_to)

    raw_cookies = handler.headers.get('X-Cookies', 'a=b')
    for cookie in raw_cookies.split(', '):
        handler.send_header('Set-Cookie', cookie)
    handler.end_headers()


@contextmanager
def _http_server():
    server = HTTPServer(('localhost', 0), TestHandler)
    thread = threading.Thread(target=server.serve_forever)
    thread.start()
    yield server
    server.shutdown()
    thread.join()


@pytest.fixture(scope="function")
def http_server():
    """A custom HTTP server implementation for our tests, that is
    built on top of the http.server module. Handy when we need to
    deal with details which httpbin can not capture."""

    with _http_server() as server:
        yield '{0}:{1}'.format(*server.socket.getsockname())


@pytest.fixture(scope="function")
def localhost_http_server():
    """Just like the http_server, but uses the static
    `localhost` name for the host."""

    with _http_server() as server:
        yield 'localhost:{1}'.format(*server.socket.getsockname())