diff --git a/bitelab/__init__.py b/bitelab/__init__.py index 8c426f3..f38def1 100644 --- a/bitelab/__init__.py +++ b/bitelab/__init__.py @@ -1,7 +1,7 @@ from typing import Optional from functools import lru_cache, wraps -from fastapi import APIRouter, Depends, FastAPI, Request +from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request from fastapi.security import OAuth2PasswordBearer from httpx import AsyncClient, Auth from starlette.status import HTTP_200_OK, HTTP_404_NOT_FOUND, HTTP_401_UNAUTHORIZED @@ -11,6 +11,7 @@ from . import data import asyncio import gc +import orm import socket import sys import tempfile @@ -103,9 +104,9 @@ oauth2_scheme = OAuth2PasswordBearer(tokenUrl='/nonexistent') async def lookup_user(token: str = Depends(oauth2_scheme), data: data.DataWrapper = Depends(get_data)): try: return (await data.APIKey.objects.get(key=token)).user - except KeyError: + except orm.exceptions.NoMatch: raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, + status_code=HTTP_401_UNAUTHORIZED, detail='Invalid authentication credentials', headers={'WWW-Authenticate': 'Bearer'}, ) @@ -213,6 +214,9 @@ class TestBiteLab(unittest.IsolatedAsyncioTestCase): res = await self.client.get('/board_info') self.assertEqual(res.status_code, HTTP_401_UNAUTHORIZED) + res = await self.client.get('/board_classes', auth=BiteAuth('badapikey')) + self.assertEqual(res.status_code, HTTP_401_UNAUTHORIZED) + async def test_classes(self): res = await self.client.get('/board_classes', auth=BiteAuth('thisisanapikey')) self.assertEqual(res.status_code, HTTP_200_OK)