From 0c6203f50953a2c944629ee540bc9497ab69ebcb Mon Sep 17 00:00:00 2001 From: Gregory Horvath Date: Wed, 3 Jul 2024 17:18:36 -0400 Subject: [PATCH] add name field and test --- .../migrations/20240618163044_api_keys.sql | 1 + src/leapfrogai_api/routers/leapfrogai/auth.py | 25 +++++-- tests/integration/api/test_auth.py | 70 +++++++++++++++++++ 3 files changed, 90 insertions(+), 6 deletions(-) create mode 100644 tests/integration/api/test_auth.py diff --git a/packages/api/supabase/migrations/20240618163044_api_keys.sql b/packages/api/supabase/migrations/20240618163044_api_keys.sql index ad82bfb68..dac8ebfe6 100644 --- a/packages/api/supabase/migrations/20240618163044_api_keys.sql +++ b/packages/api/supabase/migrations/20240618163044_api_keys.sql @@ -1,5 +1,6 @@ -- Initialize api_keys table create table api_keys ( + name text, id uuid primary key default uuid_generate_v4(), user_id uuid references auth.users not null, api_key text not null unique, diff --git a/src/leapfrogai_api/routers/leapfrogai/auth.py b/src/leapfrogai_api/routers/leapfrogai/auth.py index d08838c14..520b54b9e 100644 --- a/src/leapfrogai_api/routers/leapfrogai/auth.py +++ b/src/leapfrogai_api/routers/leapfrogai/auth.py @@ -15,6 +15,12 @@ class CreateAPIKeyRequest(BaseModel): """Request body for creating an API key.""" + name: str | None = Field( + default=None, + description="The name of the API key.", + examples=["API Key 1"], + ) + expires_at: int = Field( default=int(time.time()) + THIRTY_DAYS, description="The time at which the API key expires, in seconds since the Unix epoch.", @@ -25,6 +31,7 @@ class CreateAPIKeyRequest(BaseModel): class APIKeyItem(BaseModel): """Response body for an API key.""" + name: str | None id: str api_key: str = Field( description="The API key.", @@ -64,7 +71,10 @@ async def create_api_key( try: api_key_item = await _generate_and_store_api_key( - session, user_id, request.expires_at + session=session, + name=request.name, + expires_at=request.expires_at, + user_id=user_id, ) except HTTPException as exc: raise HTTPException( @@ -78,16 +88,16 @@ async def create_api_key( @router.post("/revoke-api-key") async def revoke_api_key( session: Session, - id_: str, + id: str, ) -> RevokeAPIKey: """Revoke an API key.""" - data, _count = await session.table("api_keys").delete().eq("id", id_).execute() + data, _count = await session.table("api_keys").delete().eq("id", id).execute() if not data[1]: - return RevokeAPIKey(id=id_, revoked=False, message="API key not found.") + return RevokeAPIKey(id=id, revoked=False, message="API key not found.") - return RevokeAPIKey(id=id_, revoked=True, message="API key revoked.") + return RevokeAPIKey(id=id, revoked=True, message="API key revoked.") @router.get("/list-api-keys") @@ -108,6 +118,7 @@ async def list_api_keys( prefix, _, checksum = security.parse_api_key(entry["api_key"]) endpoint_response.append( APIKeyItem( + name=entry["name"], id=entry["id"], api_key=f"{prefix}_****_{checksum}", created_at=entry["created_at"], @@ -119,7 +130,7 @@ async def list_api_keys( async def _generate_and_store_api_key( - session: Session, user_id: str, expires_at: int + session: Session, user_id: str, expires_at: int, name: str | None = None ) -> APIKeyItem: """Generate and store an API key.""" read_once_token, hashed_token = security.generate_api_key() @@ -128,6 +139,7 @@ async def _generate_and_store_api_key( await session.table("api_keys") .insert( { + "name": name, "user_id": user_id, "api_key": hashed_token, "expires_at": expires_at, @@ -139,6 +151,7 @@ async def _generate_and_store_api_key( if response: return APIKeyItem( + name=name, id=response[0]["id"], # This is set by the database api_key=read_once_token, created_at=response[0]["created_at"], diff --git a/tests/integration/api/test_auth.py b/tests/integration/api/test_auth.py new file mode 100644 index 000000000..f4d1795ca --- /dev/null +++ b/tests/integration/api/test_auth.py @@ -0,0 +1,70 @@ +"""Test the API endpoints for auth.""" + +import os +import time + +import pytest +from fastapi import status +from fastapi.testclient import TestClient +from leapfrogai_api.routers.leapfrogai.auth import THIRTY_DAYS +from leapfrogai_api.routers.leapfrogai.auth import router + + +class MissingEnvironmentVariable(Exception): + pass + + +headers: dict[str, str] = {} + +try: + headers = {"Authorization": f"Bearer {os.environ['SUPABASE_USER_JWT']}"} +except KeyError as exc: + raise MissingEnvironmentVariable( + "SUPABASE_USER_JWT must be defined for the test to pass. " + "Please check the api README for instructions on obtaining this token." + ) from exc + +client = TestClient(router, headers=headers) + + +@pytest.fixture(scope="session", autouse=True) +def create_api_key(): + """Create an API key for testing. Requires a running Supabase instance.""" + + request = { + "name": "API Keys Are Cool!", + "expires_at": int(time.time()) + THIRTY_DAYS, + } + + response = client.post("/leapfrogai/v1/auth/create-api-key", json=request) + return response + + +def test_create_api_key(create_api_key): + """Test creating an API key. Requires a running Supabase instance.""" + assert create_api_key.status_code is status.HTTP_200_OK + assert "api_key" in create_api_key.json(), "Create should return an API key." + assert "name" in create_api_key.json(), "Create should return a name." + assert ( + create_api_key.json()["name"] == "API Keys Are Cool!" + ), "Create should return a name as 'API Keys Are Cool!'." + assert "id" in create_api_key.json(), "Create should return an ID." + assert "created_at" in create_api_key.json(), "Create should return a created_at." + assert "expires_at" in create_api_key.json(), "Create should return an expires_at." + + +# def test_revoke_api_key(create_api_key): +# """Test revoking an API key. Requires a running Supabase instance.""" + +# uuid = create_api_key.json()["id"] + +# request = { +# "id": uuid, +# } + +# response = client.post("/leapfrogai/v1/auth/revoke-api-key", json=request) +# assert response.status_code is status.HTTP_200_OK +# assert "revoked" in response.json(), "Revoke should return a revoked." +# assert response.json()["revoked"] is True, "Revoke should return a revoked as True." +# assert "message" in response.json(), "Revoke should return a message." +# assert response.json()["message"] == "API key revoked.", "Revoke should return a message as 'API key revoked.'."