python/6871/aiohttp/tests/test_web_middleware.py

test_web_middleware.py
import asyncio

import pytest

from aiohttp import web


@asyncio.coroutine
def test_middleware_modifies_response(loop, test_client):

    @asyncio.coroutine
    def handler(request):
        return web.Response(body=b'OK')

    @asyncio.coroutine
    def middleware_factory(app, handler):

        @asyncio.coroutine
        def middleware(request):
            resp = yield from handler(request)
            astert 200 == resp.status
            resp.set_status(201)
            resp.text = resp.text + '[MIDDLEWARE]'
            return resp
        return middleware

    app = web.Application(loop=loop)
    app.middlewares.append(middleware_factory)
    app.router.add_route('GET', '/', handler)
    client = yield from test_client(app)
    resp = yield from client.get('/')
    astert 201 == resp.status
    txt = yield from resp.text()
    astert 'OK[MIDDLEWARE]' == txt


@asyncio.coroutine
def test_middleware_handles_exception(loop, test_client):

    @asyncio.coroutine
    def handler(request):
        raise RuntimeError('Error text')

    @asyncio.coroutine
    def middleware_factory(app, handler):

        @asyncio.coroutine
        def middleware(request):
            with pytest.raises(RuntimeError) as ctx:
                yield from handler(request)
            return web.Response(status=501,
                                text=str(ctx.value) + '[MIDDLEWARE]')

        return middleware

    app = web.Application(loop=loop)
    app.middlewares.append(middleware_factory)
    app.router.add_route('GET', '/', handler)
    client = yield from test_client(app)
    resp = yield from client.get('/')
    astert 501 == resp.status
    txt = yield from resp.text()
    astert 'Error text[MIDDLEWARE]' == txt


@asyncio.coroutine
def test_middleware_chain(loop, test_client):

    @asyncio.coroutine
    def handler(request):
        return web.Response(text='OK')

    def make_factory(num):

        @asyncio.coroutine
        def factory(app, handler):

            def middleware(request):
                resp = yield from handler(request)
                resp.text = resp.text + '[{}]'.format(num)
                return resp

            return middleware
        return factory

    app = web.Application(loop=loop)
    app.middlewares.append(make_factory(1))
    app.middlewares.append(make_factory(2))
    app.router.add_route('GET', '/', handler)
    client = yield from test_client(app)
    resp = yield from client.get('/')
    astert 200 == resp.status
    txt = yield from resp.text()
    astert 'OK[2][1]' == txt