I have added a middleware to my flask rest api app to add a specific header to all responses, using the after_request() decorator. What would be a good way to ensure that all endpoints include this header? I have tests for every endpoint to test the status and data of the response. I could add an extra assert in every single test to check the header too? It is of course possible that I forget to add an assert for a certain endpoint, but I dont know of a better way to test this? Any suggestions?

CodePudding user response：

I think the best way is to create a separate test, calculate registered routes and check only header and response statutes. Here is an example:

# app.py
import random
from flask import Flask, jsonify

app = Flask(__name__)


# a few routes for demo
@app.route('/user/<user_id>', methods=['GET'])
def get_user(user_id):
    return jsonify(dict(user_id=user_id))


@app.route('/user', methods=['POST'])
def create_user():
    return jsonify(dict(user_id=random.randint(0, 100000))), 201


@app.after_request
def after_request_func(response):
    # just an example - custom header
    response.headers['MY_CUSTOM_HEADER'] = 'value'
    return response

test.py:

import unittest
from parameterized import parameterized

from app import app


def get_routes_params() -> list:
    # you can move routes_map(config) to yaml and parse config before tests...
    # the test will be failed if you registered a new route(or method) and didn't set parameters
    routes_map = {
        '/user/<user_id>': {
            'GET': (dict(user_id=1), None, 200),
        },
        '/user': {
            'POST': (dict(), dict(name='Baz'), 201),
        }
    }

    params = []
    # search parameters for all registered routes from our routes_map(config)
    for rule in app.url_map.iter_rules():
        if rule.rule.startswith('/static/'):
            continue

        for method in rule.methods:
            if method in ('HEAD', 'OPTIONS'):
                continue

            route_args, json, expected_status = routes_map[rule.rule][method]
            url = rule.rule
            # replace positional route args
            for key, value in route_args.items():
                url = url.replace(f'<{key}>', str(value), 1)

            params.append([url, method.lower(), json, expected_status])
    # params for each route: [['/user/1', 'get', None, 200], ['/user', 'post', {'name': 'Baz'}, 201]]
    return params


app.config.update({'TESTING': True})
class TestMyCustomHeader(unittest.TestCase):

    @parameterized.expand(get_routes_params())
    def test_after_request_my_custom_header(self, url: str, method: str, json: dict | None, expected_status: int):
        with app.test_client() as client:
            response = getattr(client, method)(url, json=json)
            self.assertEqual(response.headers['MY_CUSTOM_HEADER'], 'value')
            self.assertEqual(response.status_code, expected_status)

So in this case you'll have a failed tests if you add a new routes(or methods) because routes calculates dynamically. All what you need is just actualize config(routes_map). Other tests will only check specific user cases, data structures, responses(positive/negative) etc.