We Keep Coding

make fastapi middleware returning custom http status instead of AuthenticationError status 400

In the following example when you pass a username in the basic auth field it raise a basic 400 error, but i want to return 401 since it's related to the authentication system.
I did tried Fastapi exceptions classes but they do not raise (i presume since we are in a starlette middleware). Il also tried JSONResponse from starlette but it doesn't work either.
AuthenticationError work and raise a 400 but it's juste an empty class that inherit from Exception so no status code can be given.
Fully working example:
import base64
import binascii
import uvicorn
from fastapi import Depends, FastAPI, HTTPException, Request
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer, HTTPBasic
from starlette.authentication import AuthenticationBackend, AuthCredentials, AuthenticationError, BaseUser
from starlette.middleware import Middleware
from starlette.middleware.authentication import AuthenticationMiddleware
from starlette.responses import JSONResponse
class SimpleUserTest(BaseUser):
"""
user object returned to route
"""
def __init__(self, username: str, test1: str, test2: str) -> None:
self.username = username
self.test1 = test1
self.test2 = test2
#property
def is_authenticated(self) -> bool:
return True
async def jwt_auth(auth: HTTPAuthorizationCredentials = Depends(HTTPBearer(auto_error=False))):
if auth:
return True
async def key_auth(apikey_header=Depends(HTTPBasic(auto_error=False))):
if apikey_header:
return True
class BasicAuthBackend(AuthenticationBackend):
async def authenticate(self, conn):
if "Authorization" not in conn.headers:
return
auth = conn.headers["Authorization"]
try:
scheme, credentials = auth.split()
if scheme.lower() == 'bearer':
# check bearer content and decode it
user: dict = {"username": "bearer", "test1": "test1", "test2": "test2"}
elif scheme.lower() == 'basic':
decoded = base64.b64decode(credentials).decode("ascii")
username, _, password = decoded.partition(":")
if username:
# check redis here instead of return dict
print("try error raise")
raise AuthenticationError('Invalid basic auth credentials') # <= raise 400, we need 401
# user: dict = {"username": "basic auth", "test1": "test1", "test2": "test2"}
else:
print("error should raise")
return JSONResponse(status_code=401, content={'reason': str("You need to provide a username")})
else:
return JSONResponse(status_code=401, content={'reason': str("Authentication type is not supported")})
except (ValueError, UnicodeDecodeError, binascii.Error) as exc:
raise AuthenticationError('Invalid basic auth credentials')
return AuthCredentials(["authenticated"]), SimpleUserTest(**user)
async def jwt_or_key_auth(jwt_result=Depends(jwt_auth), key_result=Depends(key_auth)):
if not (key_result or jwt_result):
raise HTTPException(status_code=401, detail="Not authenticated")
app = FastAPI(
dependencies=[Depends(jwt_or_key_auth)],
middleware=[Middleware(AuthenticationMiddleware, backend=BasicAuthBackend())]
)
#app.get("/")
async def read_items(request: Request) -> str:
return request.user.__dict__
if __name__ == "__main__":
uvicorn.run("main:app", host="127.0.0.1", port=5000, log_level="info")
if we set username in basic auth:
INFO: 127.0.0.1:22930 - "GET / HTTP/1.1" 400 Bad Request

so i ended up using on_error as suggested by #MatsLindh
old app:
app = FastAPI(
dependencies=[Depends(jwt_or_key_auth)],
middleware=[
Middleware(
AuthenticationMiddleware,
backend=BasicAuthBackend(),
)
],
)
new version:
app = FastAPI(
dependencies=[Depends(jwt_or_key_auth)],
middleware=[
Middleware(
AuthenticationMiddleware,
backend=BasicAuthBackend(),
on_error=lambda conn, exc: JSONResponse({"detail": str(exc)}, status_code=401),
)
],
)
I choose to use JSONResponse and return a "detail" key/value to emulate a classic 401 fastapi httperror

Python unit test coverage using pytest

I am new to pytest and wanted to add the below 3 methods for unit test coverage without actually using a real mongo db instance but rather mock it.
Could try using a real db instance but it isn't recommended.
Request for an example on how to mock mongodb client and get a document
import os
import logging
import urllib.parse
from dotenv import load_dotenv
from pymongo import MongoClient
from logger import *
load_dotenv()
def getMongoConnection():
userName = urllib.parse.quote_plus(os.getenv("USER_NAME"))
password = urllib.parse.quote_plus(os.getenv("PASSWORD"))
hostName1_port = os.getenv("HOST_NAME1")
hostName2_port = os.getenv("HOST_NAME2")
hostName3_port = os.getenv("HOST_NAME3")
authSourceDatabase = os.getenv("AUTH_SOURCE_DATABASE")
replicaSet = os.getenv("REPLICA_SET")
connectTimeoutMS = "1000"
socketTimeoutMS = "30000"
maxPoolSize = "100"
try:
client = MongoClient('mongodb://'+userName+':'+password+'#'+hostName1_port+','+hostName2_port+','+hostName3_port+'/'+authSourceDatabase+'?ssl=true&replicaSet='+replicaSet +
'&authSource='+authSourceDatabase+'&retryWrites=true&w=majority&connectTimeoutMS='+connectTimeoutMS+'&socketTimeoutMS='+socketTimeoutMS+'&maxPoolSize='+maxPoolSize)
return client
except Exception as e:
logging.error("Error while connecting to mongoDB.")
return False
def connectToDBCollection(client, databaseName, collectionName):
db = client[databaseName]
collection = db[collectionName]
return collection
def getDoc(bucketName, databaseName, collectionName):
try:
client = getMongoConnection()
if client != False:
collection = connectToDBCollection(
client, databaseName, collectionName)
return collection.find_one({'bucket': bucketName})
except Exception as e:
logging.error("An exception occurred while fetching doc, error is ", e)
Edit : (Tried using below code and was able to cover most of the cases but seeing an error)
def test_mongo():
db_conn = mongomock.MongoClient()
assert isinstance(getMongoConnection(), MongoClient)
def test_connect_mongo():
return connectToDBCollection(mongomock.MongoClient(), "sampleDB", "sampleCollection")
//trying to cover exception block for getMongoConnection()
def test_exception():
with pytest.raises(Exception) as excinfo:
getMongoConnection()
assert str(excinfo.value) == False
def test_getDoc():
collection = mongomock.MongoClient().db.collection
stored_obj = collection.find_one({'_id': 1})
assert stored_obj == getDoc("bucket", "db", "collection")
def test_createDoc():
collection = mongomock.MongoClient().db.collection
stored_obj = collection.insert_one({'_id': 1})
assert stored_obj == createDoc("bucket", "db", "collection")
def test_updateDoc():
collection = mongomock.MongoClient().db.collection
stored_obj = collection.replace_one({'_id': 1}, {'_id': 2})
assert stored_obj == updateDoc(
{'_id': 1}, {'$set': {'_id': 2}}, "db", "collection")
Errors :
test_exception - Failed: DID NOT RAISE <class 'Exception'>
test_createDoc - TypeError: not all arguments converted during string formatting
AssertionError: assert <pymongo.results.UpdateResult object at 0x7fc0e835a400> == <pymongo.results.UpdateResult object at 0x7fc0e8211900>

Looks like MongoClient is a nested dict with databaseName and collectionName or implemented with a key accessor.
You could mock the client first with
import unittest
mocked_collection = unittest.mock.MagicMock()
# mock the find_one method
mocked_collection.find_one.return_value = {'data': 'collection_find_one_result'}
mocked_client = unittest.mock.patch('pymongo.MongoClient').start()
mocked_client.return_value = {
'databaseName': {'collectionname': mocked_collection}
}

Maybe try a specialized mocking library like MongoMock?
In particular the last example using #mongomock.patch looks like it can be relevant for your code.

How to write testcase to verify 500 response for flask apis

I am trying to write unit test cases for flask APIs. I am unable to figure out how to write test cases for 500 internal server errors and custom exceptions in API response.
Here are my books api to get the book by id
#BOOKS_BP.route('/<book_id>', methods=['GET'])
def get_book(book_id):
try:
book = Book.query.filter(Book.id == book_id).scalar()
if not book:
raise DBRecordNotFound('No Record Found with ID - {}'.format(book_id))
book_schema = BookSchema()
return jsonify({'status': 'success', 'book': book_schema.dump(book)}), HTTPStatus.OK
except DBRecordNotFound as ex:
return jsonify({'status': 'error', 'message': str(ex)}), HTTPStatus.NOT_FOUND
except Exception as ex:
logger.error(ex, exc_info=True)
return jsonify({'status': 'error', 'message': str(ex)}), HTTPStatus.INTERNAL_SERVER_ERROR
Here is my get book by id test case.
import json
from tests.test_base import BaseTestCase
from tests.factories import BookFactory
from http import HTTPStatus
class BooksTestCase(BaseTestCase):
#classmethod
def setUpClass(cls):
super(BooksTestCase, cls).setUpClass()
book = BookFactory(name='Test Book')
book.save()
cls.test_book_id = book.id
def test_book_by_id(self):
# For 200 Status
response = self.client.get('/api/books/{}'.format(self.test_book_id))
data = json.loads(response.data)
assert response.status_code == HTTPStatus.OK
assert data['book']['name'] == 'Test Book'
# For 404 Status
response = self.client.get('/api/books/123456')
data = json.loads(response.data)
assert response.status_code == HTTPStatus.NOT_FOUND
assert data['message'] == 'No Record Found with ID - 123456'
# For 500 Status

status_code from pytest is not equal to the real one

Here is a fastapi webapp. I got some trouble in test.
post in crud.py:
async def post(payload: BookSchema):
query = books.insert().values(book=payload.book, author=payload.author,
create_time=payload.create_time, update_time=payload.update_time)
return await database.execute(query=query)
create_book:
#router.post("/book", response_model=BookDB, status_code=201)
async def create_book(payload: BookSchema):
book_id = await crud.post(payload)
response_object = {
"id": book_id,
**payload.dict()
}
return response_object
conftest.py:
import pytest
from starlette.testclient import TestClient
from app.main import app
#pytest.fixture()
def test_app():
with TestClient(app) as client:
yield client
I use pytest to test my case:
def test_create_book_invalid_json(test_app):
response = test_app.post(API_PREFIX + "/book", data=json.dumps({"book": "smart"}))
print(f"response.status_code is {response.status_code}")
assert response.status_code == 422
Expected status_code should be 422 and I got 201. I try to print response.status_code, and the output is 422.
test_books.py::test_create_book_invalid_json FAILED [100%]response.status_code is 422
test_books.py:20 (test_create_book_invalid_json)
201 != 422
Expected :422
Actual :201

You have not provided information about your test client here, but I assume it is a requests-like one. In such case the data parameter expects a dictionary, not a string.
response = test_app.post(url, data={'book': 'smart'})
This will send corresponding form data to your web app. If what you need is sending JSON, use the json parameter instead:
response = test_app.post(url, json={'book': 'smart'})

How to set a request's limit to a joken JWT

I'm working in an API with python, flask and implementing JWT with a timeout for expiration, but I'd like to set also a limit request, So the token is gonna be invalid if the time is out or the token has been used in five requests.
I'd been working with the timeout for expiration, but I can't find how to implement the expiration by five requests. Thanks by the help.
The Code til now:
from flask import *
import jwt
import datetime
from flask_pymongo import PyMongo
from functools import wraps
import hashlib
app = Flask(__name__)
app.config['MONGO_DBNAME'] = 'MONGOCONEX'
app.config['MONGO_URI'] = 'mongodb://localhost:27017/MONGOCONEX'
app.config['log_log_1'] = 'LOGKEYCONNECT'
app.config['key1'] = 'SECRECTKEY'
app.config['key2'] = 'PASSKEY'
mongo = PyMongo(app)
def token_required(f):
#wraps(f)
def decorated(*args, **kwargs):
token = request.args.get('token')
if not token:
return jsonify({'error': 402,'message':'Token is missing'})
try:
data = jwt.decode(token, app.config['key1'])
except:
return jsonify({'error': 403,'message': 'Token Invalid'})
return f(*args, **kwargs)
return decorated
#app.route('/results', methods=['GET'])
#token_required
def get_all_stars():
results = mongo.db.resultados
output = []
date_start = datetime.datetime.now() - datetime.timedelta(days=1*365)
date_end = datetime.datetime.now() + datetime.timedelta(days=1*365)
for s in results.find():
#print(s)
if date_start <= s['day'] <= date_end:
output.append({'day':s['day'], 'monthly_prediction':s['monthly_prediction'], 'percent_prediction':s['percent_prediction']})
return jsonify({'result' : output})
#app.route('/login', methods=['GET'])
def login():
log_key = request.args.get('l_k')
password_k = request.args.get('p_k')
md5_hash = hashlib.md5()
md5_hash.update(b""+app.config['key2']+"")
encoded_pass_key = md5_hash.hexdigest()
if (log_key == app.config['log_log_1']) and (password_k == encoded_pass_key):
token = jwt.encode({'exp': datetime.datetime.utcnow() + datetime.timedelta(minutes=2)}, app.config['key1'])
return jsonify({'token': token.decode('UTF-8')})
return jsonify({'error': 401, 'description': 'Not verified', 'Wrong Auth': 'Auth Required'})
if __name__ == '__main__':
try:
app.run(debug=True)
except Exception as e:
print('Error: '+str(e))

I see you're using mongo, the workflow is you can put counter along with the token in mongo database and count it how many it has been used, then adding logic to compare which one comes first, the time limit or how many times the token has been used, if it has been used five times you can revoke the token & generate a new token or another workflow you want to do. Here's the further reference to revoke/blacklist the token after it has been accesed five times https://flask-jwt-extended.readthedocs.io/en/stable/blacklist_and_token_revoking/