Skip to content

Commit

Permalink
fix: made token validation async!
Browse files Browse the repository at this point in the history
  • Loading branch information
amindadgar committed Nov 12, 2024
1 parent b0dd8bb commit 1515c02
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 15 deletions.
23 changes: 21 additions & 2 deletions services/api_key.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,14 +10,33 @@


async def get_api_key(api_key_header: str = Security(api_key_header)):
"""
Dependency function to validate API key
Parameters
-------------
api_key_header : str
the api key passed to the header
Raises
------
HTTPException
If API key is missing or invalid
Returns
-------
api_key_header : str
The validated API key
"""
validator = ValidateAPIKey()

if not api_key_header:
raise HTTPException(
status_code=HTTP_401_UNAUTHORIZED, detail="No API key provided"
)

if not validator.validate(api_key_header):
valid = await validator.validate(api_key_header)
if not valid:
raise HTTPException(status_code=HTTP_401_UNAUTHORIZED, detail="Invalid API key")

return api_key_header
Expand All @@ -29,7 +48,7 @@ def __init__(self) -> None:
self.db = "hivemind"
self.tokens_collection = "tokens"

def validate(self, api_key: str) -> bool:
async def validate(self, api_key: str) -> bool:
"""
check if the api key is available in mongodb or not
Expand Down
59 changes: 46 additions & 13 deletions tests/integration/test_validate_token.py
Original file line number Diff line number Diff line change
@@ -1,30 +1,49 @@
from unittest import TestCase

from unittest import IsolatedAsyncioTestCase
from services.api_key import ValidateAPIKey
from utils.mongo import MongoSingleton


class TestValidateToken(TestCase):
def setUp(self) -> None:
class TestValidateToken(IsolatedAsyncioTestCase):
async def asyncSetUp(self) -> None:
"""
Set up test case with a test database
"""
self.client = MongoSingleton.get_instance().get_client()
self.validator = ValidateAPIKey()

# changing the db so not to overlap with the right ones
# Using test database to avoid affecting production data
self.validator.db = "hivemind_test"
self.validator.tokens_collection = "tokens_test"

self.client.drop_database(self.validator.db)
# Clean start for each test
self.clean_database()

async def asyncTearDown(self) -> None:
"""
Clean up test database after each test
"""
self.clean_database()

def tearDown(self) -> None:
def clean_database(self) -> None:
"""
Helper method to clean the test database
"""
self.client.drop_database(self.validator.db)

def test_no_token_available(self):
async def test_no_token_available(self):
"""
Test validation when no tokens exist in database
"""
api_key = "1234"
valid = self.validator.validate(api_key)
valid = await self.validator.validate(api_key)

self.assertEqual(valid, False)

def test_no_matching_token_available(self):
async def test_no_matching_token_available(self):
"""
Test validation when tokens exist but none match
"""
# Insert test tokens - no await needed as this is synchronous
self.client[self.validator.db][self.validator.tokens_collection].insert_many(
[
{
Expand All @@ -44,13 +63,18 @@ def test_no_matching_token_available(self):
},
]
)

api_key = "1234"
valid = self.validator.validate(api_key)
valid = await self.validator.validate(api_key)

self.assertEqual(valid, False)

def test_single_token_available(self):
async def test_single_token_available(self):
"""
Test validation when matching token exists
"""
api_key = "1234"

self.client[self.validator.db][self.validator.tokens_collection].insert_many(
[
{
Expand All @@ -70,6 +94,15 @@ def test_single_token_available(self):
},
]
)
valid = self.validator.validate(api_key)

valid = await self.validator.validate(api_key)

self.assertEqual(valid, True)

async def test_validation_with_empty_api_key(self):
"""
Test validation with empty API key
"""
valid = await self.validator.validate("")

self.assertEqual(valid, False)

0 comments on commit 1515c02

Please sign in to comment.