feat(agents-api): Add route tests

agents-api/agents_api/clients/cozo.py

@@ -1,9 +1,16 @@
 from pycozo.client import Client
 
 from ..env import cozo_auth, cozo_host
+from ..web import app
 
 options = {"host": cozo_host}
 if cozo_auth:
     options.update({"auth": cozo_auth})
 
-client = Client("http", options=options)
+
+def get_cozo_client():
+    client = getattr(app.state, "cozo_client", Client("http", options=options))
+    if not hasattr(app.state, "cozo_client"):
+        app.state.cozo_client = client
+
+    return client

agents-api/agents_api/env.py

@@ -3,13 +3,13 @@
 It utilizes the environs library for environment variable parsing.
 """
 
+import random
 from pprint import pprint
 
 from environs import Env
 
 # Initialize the Env object for environment variable parsing.
 env = Env()
-env.read_env()
 
 # Debug mode
 debug: bool = env.bool("AGENTS_API_DEBUG", default=False)
@@ -32,7 +32,12 @@
 temporal_task_queue = env.str("TEMPORAL_TASK_QUEUE", default="memory-task-queue")
 
 # auth
-api_key: str = env.str("AGENTS_API_KEY")
+_random_generated_key = "".join(str(random.randint(0, 9)) for _ in range(32))
+api_key: str = env.str("AGENTS_API_KEY", _random_generated_key)
+
+if api_key == _random_generated_key:
+    print(f"Generated API key since not set in the environment: {api_key}")
+
 api_key_header_name: str = env.str("AGENTS_API_KEY_HEADER_NAME", default="X-Auth-Key")
 skip_check_developer_headers: bool = env.bool(
     "SKIP_CHECK_DEVELOPER_HEADERS", default=False

agents-api/agents_api/models/agent/get_agent.py

@@ -18,10 +18,14 @@
 
 @rewrap_exceptions(
     {
+        lambda e: isinstance(e, QueryException)
+        and "Developer not found" in str(e): lambda *_: HTTPException(
+            detail="developer does not exist", status_code=403
+        ),
         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
+            detail="developer doesnt own resource", status_code=404
         ),
         QueryException: partialclass(HTTPException, status_code=400),
         ValidationError: partialclass(HTTPException, status_code=400),

agents-api/agents_api/models/user/create_user.py

@@ -22,6 +22,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),

agents-api/agents_api/models/user/delete_user.py

@@ -57,15 +57,18 @@ def delete_user(*, developer_id: UUID, user_id: UUID) -> tuple[list[str], dict]:
         verify_developer_owns_resource_query(developer_id, "users", user_id=user_id),
         """
         # Delete docs
-        ?[user_id, doc_id] :=
+        ?[owner_type, owner_id, doc_id] :=
             *docs{
-                owner_id: user_id,
-                owner_type: "user",
+                owner_id,
+                owner_type,
                 doc_id,
-            }, user_id = to_uuid($user_id)
+            },
+            owner_id = to_uuid($user_id),
+            owner_type = "user"
 
         :delete docs {
-            user_id,
+            owner_type,
+            owner_id,
             doc_id
         }
         :returning

agents-api/agents_api/models/user/get_user.py

@@ -18,6 +18,15 @@
 
 @rewrap_exceptions(
     {
+        lambda e: isinstance(e, QueryException)
+        and "Developer not found" in str(e): lambda *_: HTTPException(
+            detail="developer does not exist", status_code=403
+        ),
+        lambda e: isinstance(e, QueryException)
+        and "asserted to return some results, but returned none"
+        in str(e): lambda *_: HTTPException(
+            detail="developer doesnt own resource", status_code=404
+        ),
         QueryException: partialclass(HTTPException, status_code=400),
         ValidationError: partialclass(HTTPException, status_code=400),
         TypeError: partialclass(HTTPException, status_code=400),

agents-api/agents_api/models/utils.py

@@ -7,7 +7,6 @@
 import pandas as pd
 from pydantic import BaseModel
 
-from ..clients.cozo import client as cozo_client
 from ..common.utils.cozo import uuid_int_list_to_uuid4
 
 P = ParamSpec("P")
@@ -66,12 +65,15 @@ class NewCls(cls):
 
 def verify_developer_id_query(developer_id: UUID | str) -> str:
     return f"""
-    ?[developer_id] :=
+    matched[count(developer_id)] :=
         *developers{{
             developer_id,
         }}, developer_id = to_uuid("{str(developer_id)}")
 
-    :assert some
+    ?[exists] :=
+        matched[num],
+        exists = num > 0,
+        assert(exists, "Developer does not exist")
     """
 
 
@@ -127,9 +129,7 @@ def cozo_query_dec(func: Callable[P, tuple[str | list[Any], dict]]):
             from pprint import pprint
 
         @wraps(func)
-        def wrapper(
-            *args: P.args, client=cozo_client, **kwargs: P.kwargs
-        ) -> pd.DataFrame:
+        def wrapper(*args: P.args, client=None, **kwargs: P.kwargs) -> pd.DataFrame:
             queries, variables = func(*args, **kwargs)
 
             if isinstance(queries, str):
@@ -146,6 +146,9 @@ def wrapper(
                 )
             )
 
+            from ..clients.cozo import get_cozo_client
+
+            client = client or get_cozo_client()
             result = client.run(query, variables)
 
             # Need to fix the UUIDs in the result

agents-api/agents_api/routers/agents/create_agent.py

@@ -19,9 +19,10 @@ async def create_agent(
     x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
     data: CreateAgentRequest,
 ) -> ResourceCreatedResponse:
+    print("create_agent", x_developer_id, data)
     agent = models.agent.create_agent(
         developer_id=x_developer_id,
         data=data,
     )
 
-    return ResourceCreatedResponse(id=agent.id, created_at=agent.created_at)
+    return ResourceCreatedResponse(id=agent.id, created_at=agent.created_at, jobs=[])

agents-api/agents_api/routers/agents/create_or_update_agent.py

@@ -27,4 +27,4 @@ async def create_or_update_agent(
         data=data,
     )
 
-    return ResourceCreatedResponse(id=agent.id, created_at=agent.created_at)
+    return ResourceCreatedResponse(id=agent.id, created_at=agent.created_at, jobs=[])

agents-api/agents_api/routers/users/create_user.py

@@ -20,4 +20,4 @@ async def create_user(
         data=data,
     )
 
-    return ResourceCreatedResponse(id=user.id, created_at=user.created_at)
+    return ResourceCreatedResponse(id=user.id, created_at=user.created_at, jobs=[])

agents-api/agents_api/routers/users/delete_user.py

@@ -14,4 +14,5 @@
 async def delete_user(
     user_id: UUID4, x_developer_id: Annotated[UUID4, Depends(get_developer_id)]
 ) -> ResourceDeletedResponse:
+    print(user_id)
     return delete_user_query(developer_id=x_developer_id, user_id=user_id)

agents-api/agents_api/routers/users/list_users.py

@@ -1,6 +1,6 @@
 import json
 from json import JSONDecodeError
-from typing import Annotated, List, Literal
+from typing import Annotated, Literal
 
 from fastapi import Depends, HTTPException, status
 from pydantic import UUID4
@@ -19,7 +19,7 @@ async def list_users(
     sort_by: Literal["created_at", "updated_at"] = "created_at",
     direction: Literal["asc", "desc"] = "desc",
     metadata_filter: str = "{}",
-) -> List[User]:
+) -> ListResponse[User]:
     try:
         metadata_filter = json.loads(metadata_filter)
     except JSONDecodeError:
@@ -37,4 +37,5 @@ async def list_users(
         metadata_filter=metadata_filter,
     )
 
-    return ListResponse[User](items=users)
+    result = ListResponse[User](items=users)
+    return result

agents-api/agents_api/web.py

@@ -22,6 +22,7 @@
 from agents_api.routers import (
     agents,
     tasks,
+    users,
 )
 
 if not sentry_dsn:
@@ -88,7 +89,7 @@ def register_exceptions(app: FastAPI):
 
 app.include_router(agents.router)
 # app.include_router(sessions.router)
-# app.include_router(users.router)
+app.include_router(users.router)
 # app.include_router(jobs.router)
 app.include_router(tasks.router)
 

agents-api/tests/fixtures.py

@@ -1,7 +1,7 @@
 from uuid import uuid4
 
 from cozo_migrate.api import apply, init
-from julep import AsyncClient, Client
+from fastapi.testclient import TestClient
 from pycozo import Client as CozoClient
 from temporalio.client import WorkflowHandle
 from ward import fixture
@@ -15,6 +15,7 @@
     CreateToolRequest,
     CreateUserRequest,
 )
+from agents_api.env import api_key, api_key_header_name
 from agents_api.models.agent.create_agent import create_agent
 from agents_api.models.agent.delete_agent import delete_agent
 from agents_api.models.docs.create_doc import create_doc
@@ -28,26 +29,7 @@
 from agents_api.models.tools.delete_tool import delete_tool
 from agents_api.models.user.create_user import create_user
 from agents_api.models.user.delete_user import delete_user
-
-# TODO: make clients connect to real service
-
-
-@fixture(scope="global")
-def client():
-    # Mock server base url
-    base_url = "http://localhost:8080"
-    client = Client(api_key="thisisnotarealapikey", base_url=base_url)
-
-    return client
-
-
-@fixture
-def async_client():
-    # Mock server base url
-    base_url = "http://localhost:8080"
-    client = AsyncClient(api_key="thisisnotarealapikey", base_url=base_url)
-
-    return client
+from agents_api.web import app
 
 
 @fixture(scope="global")
@@ -263,3 +245,26 @@ def test_tool(
         tool_id=tool.id,
         client=client,
     )
+
+
+@fixture(scope="global")
+def client(cozo_client=cozo_client):
+    client = TestClient(app=app)
+    app.state.cozo_client = cozo_client
+
+    return client
+
+
+@fixture(scope="global")
+def make_request(client=client, developer_id=test_developer_id):
+    def _make_request(method, url, **kwargs):
+        headers = kwargs.pop("headers", {})
+        headers = {
+            **headers,
+            "X-Developer-Id": str(developer_id),
+            api_key_header_name: api_key,
+        }
+
+        return client.request(method, url, headers=headers, **kwargs)
+
+    return _make_request