add tests for items handler and similar

Author: john0isaac

src/fastapi_app/__init__.py

@@ -18,7 +18,7 @@
 async def lifespan(app: FastAPI):
     load_dotenv(override=True)
 
-    azure_credential = None
+    azure_credential: azure.identity.DefaultAzureCredential | azure.identity.ManagedIdentityCredential | None = None
     try:
         if client_id := os.getenv("APP_IDENTITY_ID"):
             # Authenticate using a user-assigned managed identity on Azure

src/fastapi_app/api_routes.py

@@ -1,4 +1,5 @@
 import fastapi
+from fastapi import HTTPException
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import async_sessionmaker
 
@@ -20,6 +21,8 @@ async def item_handler(id: int):
     async_session_maker = async_sessionmaker(global_storage.engine, expire_on_commit=False)
     async with async_session_maker() as session:
         item = (await session.scalars(select(Item).where(Item.id == id))).first()
+        if not item:
+            raise HTTPException(detail=f"Item with ID {id} not found.", status_code=404)
         return item.to_dict()
 
 
@@ -29,6 +32,8 @@ async def similar_handler(id: int, n: int = 5):
     async_session_maker = async_sessionmaker(global_storage.engine, expire_on_commit=False)
     async with async_session_maker() as session:
         item = (await session.scalars(select(Item).where(Item.id == id))).first()
+        if not item:
+            raise HTTPException(detail=f"Item with ID {id} not found.", status_code=404)
         closest = await session.execute(
             select(Item, Item.embedding.l2_distance(item.embedding))
             .filter(Item.id != id)

src/fastapi_app/rag_advanced.py

@@ -97,9 +97,10 @@ async def run(
             n=1,
             stream=False,
         )
-        first_choice = chat_completion_response.choices[0]
+        first_choice_message = chat_completion_response.choices[0].message
+
         return RetrievalResponse(
-            message=Message(content=first_choice.message.content, role=first_choice.message.role),
+            message=Message(content=str(first_choice_message.content), role=first_choice_message.role),
             context=RAGContext(
                 data_points={item.id: item.to_dict() for item in results},
                 thoughts=[

src/fastapi_app/rag_simple.py

@@ -67,9 +67,10 @@ async def run(
             n=1,
             stream=False,
         )
-        first_choice = chat_completion_response.choices[0]
+        first_choice_message = chat_completion_response.choices[0].message
+
         return RetrievalResponse(
-            message=Message(content=first_choice.message.content, role=first_choice.message.role),
+            message=Message(content=str(first_choice_message.content), role=first_choice_message.role),
             context=RAGContext(
                 data_points={item.id: item.to_dict() for item in results},
                 thoughts=[

tests/conftest.py

@@ -93,8 +93,8 @@ async def db_session():
     engine = await create_postgres_engine_from_env()
     async_sesion = async_sessionmaker(autocommit=False, autoflush=False, bind=engine)
     session = async_sesion()
-    session.begin()
+    await session.begin()
     yield session
-    session.rollback()
-    session.close()
+    await session.rollback()
+    await session.close()
     await engine.dispose()