Merge pull request #449 from julep-ai/x/agents-api-tests

fix(agents-api): Fixed tests for models and queries
julep-ai · Aug 9, 2024 · db58349 · db58349
2 parents 31a485e + ad27ec0
commit db58349
Show file tree

Hide file tree

Showing 26 changed files with 1,101 additions and 1,593 deletions.
diff --git a/agents-api/agents_api/models/agent/create_agent.py b/agents-api/agents_api/models/agent/create_agent.py
@@ -23,6 +23,11 @@
 
 @rewrap_exceptions(
     {
+        lambda e: isinstance(e, QueryException)
+        and "asserted to return some results, but returned none"
+        in str(e): lambda *_: HTTPException(
+            detail="developer not found", status_code=403
+        ),
         QueryException: partialclass(HTTPException, status_code=400),
         ValidationError: partialclass(HTTPException, status_code=400),
         TypeError: partialclass(HTTPException, status_code=400),
@@ -62,7 +67,7 @@ def create_agent(
     )
     data.default_settings = data.default_settings or {}
 
-    agent_data = data.model_dump()
+    agent_data = data.model_dump(exclude_unset=True)
     default_settings = agent_data.pop("default_settings")
 
     settings_cols, settings_vals = cozo_process_mutate_data(

diff --git a/agents-api/agents_api/models/agent/delete_agent.py b/agents-api/agents_api/models/agent/delete_agent.py
@@ -23,6 +23,11 @@
 
 @rewrap_exceptions(
     {
+        lambda e: isinstance(e, QueryException)
+        and "asserted to return some results, but returned none"
+        in str(e): lambda *_: HTTPException(
+            detail="developer not found or doesnt own resource", status_code=404
+        ),
         QueryException: partialclass(HTTPException, status_code=400),
         ValidationError: partialclass(HTTPException, status_code=400),
         TypeError: partialclass(HTTPException, status_code=400),
@@ -57,14 +62,18 @@ def delete_agent(*, developer_id: UUID, agent_id: UUID) -> tuple[list[str], dict
         verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id),
         """
         # Delete docs
-        ?[agent_id, doc_id] :=
-            *agent_docs{
-                agent_id,
+        ?[owner_id, owner_type, doc_id] :=
+            *docs{
+                owner_type,
+                owner_id,
                 doc_id,
-            }, agent_id = to_uuid($agent_id)
+            },
+            owner_id = to_uuid($agent_id),
+            owner_type = "agent"
 
-        :delete agent_docs {
-            agent_id,
+        :delete docs {
+            owner_type,
+            owner_id,
             doc_id
         }
         :returning

diff --git a/agents-api/agents_api/models/agent/get_agent.py b/agents-api/agents_api/models/agent/get_agent.py
@@ -18,6 +18,11 @@
 
 @rewrap_exceptions(
     {
+        lambda e: isinstance(e, QueryException)
+        and "asserted to return some results, but returned none"
+        in str(e): lambda *_: HTTPException(
+            detail="developer not found or doesnt own resource", status_code=404
+        ),
         QueryException: partialclass(HTTPException, status_code=400),
         ValidationError: partialclass(HTTPException, status_code=400),
         TypeError: partialclass(HTTPException, status_code=400),

diff --git a/agents-api/agents_api/models/agent/test_agent_queries.py b/agents-api/agents_api/models/agent/test_agent_queries.py
diff --git a/agents-api/agents_api/models/docs/delete_doc.py b/agents-api/agents_api/models/docs/delete_doc.py
@@ -12,6 +12,7 @@
     partialclass,
     rewrap_exceptions,
     verify_developer_id_query,
+    verify_developer_owns_resource_query,
     wrap_in_class,
 )
 
@@ -37,6 +38,8 @@
 def delete_doc(
     *,
     developer_id: UUID,
+    owner_id: UUID,
+    owner_type: str,
     doc_id: UUID,
 ) -> tuple[list[str], dict]:
     """Constructs and returns a datalog query for deleting documents and associated information snippets.
@@ -52,6 +55,7 @@ def delete_doc(
     """
     # Convert UUID parameters to string format for use in the datalog query
     doc_id = str(doc_id)
+    owner_id = str(owner_id)
 
     # The following query is divided into two main parts:
     # 1. Deleting information snippets associated with the document
@@ -75,16 +79,19 @@ def delete_doc(
 
     delete_doc_query = """
         # Delete the docs
-        ?[doc_id] <- [[ to_uuid($doc_id) ]]
+        ?[doc_id, owner_type, owner_id] <- [[ to_uuid($doc_id), $owner_type, to_uuid($owner_id) ]]
 
-        :delete docs { doc_id }
+        :delete docs { doc_id, owner_type, owner_id }
         :returning
     """
 
     queries = [
         verify_developer_id_query(developer_id),
+        verify_developer_owns_resource_query(
+            developer_id, f"{owner_type}s", **{f"{owner_type}_id": owner_id}
+        ),
         delete_snippets_query,
         delete_doc_query,
     ]
 
-    return (queries, {"doc_id": doc_id})
+    return (queries, {"doc_id": doc_id, "owner_type": owner_type, "owner_id": owner_id})