fix(agents-api): add more coverage to agent routes

julep-ai · Dec 27, 2024 · a3c92c3 · a3c92c3
1 parent cf4c5ef
commit a3c92c3
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 6 deletions.
diff --git a/agents-api/agents_api/queries/agents/delete_agent.py b/agents-api/agents_api/queries/agents/delete_agent.py
@@ -3,6 +3,7 @@
 It constructs and executes SQL queries to remove agent records and associated data.
 """
 
+from typing import Literal
 from uuid import UUID
 
 import asyncpg
@@ -85,6 +86,11 @@
             status_code=400,
             detail="Invalid data provided. Please check the input values.",
         ),
+        asyncpg.exceptions.NoDataFoundError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="The specified agent does not exist.",
+        ),
     }
 )
 @wrap_in_class(
@@ -94,7 +100,7 @@
 )
 @pg_query
 @beartype
-async def delete_agent(*, agent_id: UUID, developer_id: UUID) -> tuple[str, list]:
+async def delete_agent(*, agent_id: UUID, developer_id: UUID) -> tuple[str, list, Literal["fetch", "fetchrow", "fetchmany"]]:
     """
     Constructs the SQL query to delete an agent and its related settings.
 
@@ -111,4 +117,5 @@ async def delete_agent(*, agent_id: UUID, developer_id: UUID) -> tuple[str, list
     return (
         agent_query,
         params,
+        "fetchrow",
     )
diff --git a/agents-api/agents_api/queries/agents/patch_agent.py b/agents-api/agents_api/queries/agents/patch_agent.py
@@ -3,6 +3,7 @@
 It constructs and executes SQL queries to update specific fields of an agent based on agent ID and developer ID.
 """
 
+from typing import Literal
 from uuid import UUID
 
 import asyncpg
@@ -70,6 +71,11 @@
             status_code=400,
             detail="Invalid data provided. Please check the input values.",
         ),
+        asyncpg.exceptions.NoDataFoundError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="The specified agent does not exist.",
+        ),
     }
 )
 @wrap_in_class(
@@ -82,7 +88,7 @@
 @beartype
 async def patch_agent(
     *, agent_id: UUID, developer_id: UUID, data: PatchAgentRequest
-) -> tuple[str, list]:
+) -> tuple[str, list, Literal["fetch", "fetchrow", "fetchmany"]]:
     """
     Constructs the SQL query to partially update an agent's details.
 
@@ -107,4 +113,5 @@ async def patch_agent(
     return (
         agent_query,
         params,
+        "fetchrow",
     )
diff --git a/agents-api/agents_api/queries/agents/update_agent.py b/agents-api/agents_api/queries/agents/update_agent.py
@@ -3,6 +3,7 @@
 It constructs and executes SQL queries to replace an agent's details based on agent ID and developer ID.
 """
 
+from typing import Literal
 from uuid import UUID
 
 import asyncpg
@@ -55,6 +56,11 @@
             status_code=400,
             detail="Invalid data provided. Please check the input values.",
         ),
+        asyncpg.exceptions.NoDataFoundError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="The specified agent does not exist.",
+        ),
     }
 )
 @wrap_in_class(
@@ -67,7 +73,7 @@
 @beartype
 async def update_agent(
     *, agent_id: UUID, developer_id: UUID, data: UpdateAgentRequest
-) -> tuple[str, list]:
+) -> tuple[str, list, Literal["fetch", "fetchrow", "fetchmany"]]:
     """
     Constructs the SQL query to fully update an agent's details.
 
@@ -92,4 +98,5 @@ async def update_agent(
     return (
         agent_query,
         params,
+        "fetchrow",
     )
diff --git a/agents-api/tests/test_agent_routes.py b/agents-api/tests/test_agent_routes.py
@@ -102,7 +102,7 @@ def _(make_request=make_request, agent=test_agent):
     assert response.status_code != 404
 
 
-@test("route: delete agent")
+@test("route: delete agent - exists")
 def _(make_request=make_request):
     data = dict(
         name="test agent",
@@ -132,8 +132,18 @@ def _(make_request=make_request):
 
     assert response.status_code == 404
 
+@test("route: delete agent - not exists")
+def _(make_request=make_request):
+    agent_id = str(uuid7())
 
-@test("route: update agent")
+    response = make_request(
+        method="DELETE",
+        url=f"/agents/{agent_id}",
+    )
+
+    assert response.status_code == 404
+
+@test("route: update agent - exists")
 def _(make_request=make_request, agent=test_agent):
     data = dict(
         name="updated agent",
@@ -164,8 +174,25 @@ def _(make_request=make_request, agent=test_agent):
 
     assert "test" not in agent["metadata"]
 
+@test("route: update agent - not exists")
+def _(make_request=make_request):
+    agent_id = str(uuid7())
 
-@test("route: patch agent")
+    data = dict(
+        name="updated agent",
+        about="updated agent about",
+        default_settings={"temperature": 1.0},
+    )
+
+    response = make_request(
+        method="PUT",
+        url=f"/agents/{agent_id}",
+        json=data,
+    )
+
+    assert response.status_code == 404
+
+@test("route: patch agent - exists")
 def _(make_request=make_request, agent=test_agent):
     agent_id = str(agent.id)
 
@@ -196,6 +223,23 @@ def _(make_request=make_request, agent=test_agent):
 
     assert "hello" in agent["metadata"]
 
+@test("route: patch agent - not exists")
+def _(make_request=make_request):
+    agent_id = str(uuid7())
+
+    data = dict(
+        name="patched agent",
+        about="patched agent about",
+        default_settings={"temperature": 1.0},
+    )
+
+    response = make_request(
+        method="PATCH",
+        url=f"/agents/{agent_id}",
+        json=data,
+    )
+
+    assert response.status_code == 404
 
 @test("route: list agents")
 def _(make_request=make_request):