added api key manager

backend/tests/integration/common_utils/managers/api_key.py

@@ -0,0 +1,92 @@
+from uuid import uuid4
+
+import requests
+
+from danswer.db.models import UserRole
+from ee.danswer.server.api_key.models import APIKeyArgs
+from tests.integration.common_utils.constants import API_SERVER_URL
+from tests.integration.common_utils.constants import GENERAL_HEADERS
+from tests.integration.common_utils.test_models import TestAPIKey
+from tests.integration.common_utils.test_models import TestUser
+
+
+class APIKeyManager:
+    @staticmethod
+    def create(
+        name: str | None = None,
+        api_key_role: UserRole = UserRole.ADMIN,
+        user_performing_action: TestUser | None = None,
+    ) -> TestAPIKey:
+        name = f"{name}-api-key" if name else f"test-api-key-{uuid4()}"
+        api_key_request = APIKeyArgs(
+            name=name,
+            role=api_key_role,
+        )
+        api_key_response = requests.post(
+            f"{API_SERVER_URL}/admin/api-key",
+            json=api_key_request.model_dump(),
+            headers=user_performing_action.headers
+            if user_performing_action
+            else GENERAL_HEADERS,
+        )
+        api_key_response.raise_for_status()
+        api_key = api_key_response.json()
+        result_api_key = TestAPIKey(
+            api_key_id=api_key["api_key_id"],
+            api_key_display=api_key["api_key_display"],
+            api_key=api_key["api_key"],
+            api_key_name=name,
+            api_key_role=api_key_role,
+            user_id=api_key["user_id"],
+            headers=GENERAL_HEADERS,
+        )
+        result_api_key.headers["Authorization"] = f"Bearer {result_api_key.api_key}"
+        return result_api_key
+
+    @staticmethod
+    def delete(
+        api_key: TestAPIKey,
+        user_performing_action: TestUser | None = None,
+    ) -> None:
+        api_key_response = requests.delete(
+            f"{API_SERVER_URL}/admin/api-key/{api_key.api_key_id}",
+            headers=user_performing_action.headers
+            if user_performing_action
+            else GENERAL_HEADERS,
+        )
+        api_key_response.raise_for_status()
+
+    @staticmethod
+    def get_all(
+        user_performing_action: TestUser | None = None,
+    ) -> list[TestAPIKey]:
+        api_key_response = requests.get(
+            f"{API_SERVER_URL}/admin/api-key",
+            headers=user_performing_action.headers
+            if user_performing_action
+            else GENERAL_HEADERS,
+        )
+        api_key_response.raise_for_status()
+        return [TestAPIKey(**api_key) for api_key in api_key_response.json()]
+
+    @staticmethod
+    def verify(
+        api_key: TestAPIKey,
+        verify_deleted: bool = False,
+        user_performing_action: TestUser | None = None,
+    ) -> None:
+        retrieved_keys = APIKeyManager.get_all(
+            user_performing_action=user_performing_action
+        )
+        for key in retrieved_keys:
+            if key.api_key_id == api_key.api_key_id:
+                if verify_deleted:
+                    raise ValueError("API Key found when it should have been deleted")
+                if (
+                    key.api_key_name == api_key.api_key_name
+                    and key.api_key_role == api_key.api_key_role
+                ):
+                    return
+
+        if not verify_deleted:
+            raise Exception("API Key not found")

backend/tests/integration/common_utils/managers/document.py

@@ -6,6 +6,7 @@
 from tests.integration.common_utils.constants import API_SERVER_URL
 from tests.integration.common_utils.constants import GENERAL_HEADERS
 from tests.integration.common_utils.constants import NUM_DOCS
+from tests.integration.common_utils.managers.api_key import TestAPIKey
 from tests.integration.common_utils.managers.cc_pair import TestCCPair
 from tests.integration.common_utils.test_models import SimpleTestDocument
 from tests.integration.common_utils.test_models import TestUser
@@ -79,7 +80,7 @@ def seed_and_attach_docs(
         cc_pair: TestCCPair,
         num_docs: int = NUM_DOCS,
         document_ids: list[str] | None = None,
-        user_with_api_key: TestUser | None = None,
+        api_key: TestAPIKey | None = None,
     ) -> TestCCPair:
         # Use provided document_ids if available, otherwise generate random UUIDs
         if document_ids is None:
@@ -94,9 +95,7 @@ def seed_and_attach_docs(
             response = requests.post(
                 f"{API_SERVER_URL}/danswer-api/ingestion",
                 json=document,
-                headers=user_with_api_key.headers
-                if user_with_api_key
-                else GENERAL_HEADERS,
+                headers=api_key.headers if api_key else GENERAL_HEADERS,
             )
             response.raise_for_status()
 

backend/tests/integration/common_utils/managers/user.py

@@ -8,7 +8,6 @@
 from danswer.server.manage.models import AllUsersResponse
 from danswer.server.models import FullUserSnapshot
 from danswer.server.models import InvitedUserSnapshot
-from ee.danswer.server.api_key.models import APIKeyArgs
 from tests.integration.common_utils.constants import API_SERVER_URL
 from tests.integration.common_utils.constants import GENERAL_HEADERS
 from tests.integration.common_utils.test_models import TestUser
@@ -74,29 +73,6 @@ def login_as_user(test_user: TestUser) -> str:
         print(f"Logged in as {test_user.email}")
         return f"{result_cookie.name}={result_cookie.value}"
 
-    @staticmethod
-    def add_api_key_to_user(
-        user: TestUser,
-        name: str | None = None,
-        api_key_role: UserRole = UserRole.ADMIN,
-        user_performing_action: TestUser | None = None,
-    ) -> TestUser:
-        if user_performing_action is None:
-            user_performing_action = user
-        name = f"{name}-api-key" if name else f"test-api-key-{uuid4()}"
-        api_key_request = APIKeyArgs(
-            name=name,
-            role=api_key_role,
-        )
-        api_key_response = requests.post(
-            f"{API_SERVER_URL}/admin/api-key",
-            json=api_key_request.model_dump(),
-            headers=user_performing_action.headers,
-        )
-        api_key_response.raise_for_status()
-        user.headers["Authorization"] = f"Bearer {api_key_response.json()['api_key']}"
-        return user
-
     @staticmethod
     def verify_role(user_to_verify: TestUser, target_role: UserRole) -> bool:
         response = requests.get(

backend/tests/integration/common_utils/managers/user_group.py

@@ -59,16 +59,17 @@ def edit(
         response.raise_for_status()
 
     @staticmethod
-    def set_user_to_curator(
+    def set_curator_status(
         test_user_group: TestUserGroup,
         user_to_set_as_curator: TestUser,
+        is_curator: bool = True,
         user_performing_action: TestUser | None = None,
     ) -> None:
         if not user_to_set_as_curator.id:
             raise ValueError("User has no ID")
         set_curator_request = {
             "user_id": user_to_set_as_curator.id,
-            "is_curator": True,
+            "is_curator": is_curator,
         }
         response = requests.post(
             f"{API_SERVER_URL}/manage/admin/user-group/{test_user_group.id}/set-curator",

backend/tests/integration/common_utils/test_models.py

@@ -1,8 +1,10 @@
 from typing import Any
+from uuid import UUID
 
 from pydantic import BaseModel
 from pydantic import Field
 
+from danswer.auth.schemas import UserRole
 from danswer.search.enums import RecencyBiasSetting
 from danswer.server.documents.models import DocumentSource
 from danswer.server.documents.models import InputType
@@ -17,6 +19,17 @@
 """
 
 
+class TestAPIKey(BaseModel):
+    api_key_id: int
+    api_key_display: str
+    api_key: str | None = None  # only present on initial creation
+    api_key_name: str | None = None
+    api_key_role: UserRole
+
+    user_id: UUID
+    headers: dict
+
+
 class TestUser(BaseModel):
     id: str
     email: str

backend/tests/integration/tests/connector/test_deletion.py

@@ -9,22 +9,24 @@
 
 from danswer.server.documents.models import DocumentSource
 from tests.integration.common_utils.constants import NUM_DOCS
+from tests.integration.common_utils.managers.api_key import APIKeyManager
 from tests.integration.common_utils.managers.cc_pair import CCPairManager
 from tests.integration.common_utils.managers.document import DocumentManager
 from tests.integration.common_utils.managers.document_set import DocumentSetManager
-from tests.integration.common_utils.managers.user import TestUser
 from tests.integration.common_utils.managers.user import UserManager
-from tests.integration.common_utils.managers.user_group import TestUserGroup
 from tests.integration.common_utils.managers.user_group import UserGroupManager
+from tests.integration.common_utils.test_models import TestAPIKey
+from tests.integration.common_utils.test_models import TestUser
+from tests.integration.common_utils.test_models import TestUserGroup
 from tests.integration.common_utils.vespa import TestVespaClient
 
 
 def test_connector_deletion(reset: None, vespa_client: TestVespaClient) -> None:
     # Creating an admin user (first user created is automatically an admin)
     admin_user: TestUser = UserManager.create(name="admin_user")
     # add api key to user
-    admin_user = UserManager.add_api_key_to_user(
-        user=admin_user,
+    api_key: TestAPIKey = APIKeyManager.create(
+        user_performing_action=admin_user,
     )
 
     # create connectors
@@ -41,12 +43,12 @@ def test_connector_deletion(reset: None, vespa_client: TestVespaClient) -> None:
     cc_pair_1 = DocumentManager.seed_and_attach_docs(
         cc_pair=cc_pair_1,
         num_docs=NUM_DOCS,
-        user_with_api_key=admin_user,
+        api_key=api_key,
     )
     cc_pair_2 = DocumentManager.seed_and_attach_docs(
         cc_pair=cc_pair_2,
         num_docs=NUM_DOCS,
-        user_with_api_key=admin_user,
+        api_key=api_key,
     )
 
     # create document sets
@@ -152,8 +154,8 @@ def test_connector_deletion_for_overlapping_connectors(
     # Creating an admin user (first user created is automatically an admin)
     admin_user: TestUser = UserManager.create(name="admin_user")
     # add api key to user
-    admin_user = UserManager.add_api_key_to_user(
-        user=admin_user,
+    api_key: TestAPIKey = APIKeyManager.create(
+        user_performing_action=admin_user,
     )
 
     # create connectors
@@ -170,12 +172,12 @@ def test_connector_deletion_for_overlapping_connectors(
     cc_pair_1 = DocumentManager.seed_and_attach_docs(
         cc_pair=cc_pair_1,
         document_ids=doc_ids,
-        user_with_api_key=admin_user,
+        api_key=api_key,
     )
     cc_pair_2 = DocumentManager.seed_and_attach_docs(
         cc_pair=cc_pair_2,
         document_ids=doc_ids,
-        user_with_api_key=admin_user,
+        api_key=api_key,
     )
 
     # verify vespa document exists and that it is not in any document sets or groups

backend/tests/integration/tests/dev_apis/test_simple_chat_api.py

@@ -4,11 +4,13 @@
 from tests.integration.common_utils.constants import API_SERVER_URL
 from tests.integration.common_utils.constants import NUM_DOCS
 from tests.integration.common_utils.llm import LLMProviderManager
+from tests.integration.common_utils.managers.api_key import APIKeyManager
 from tests.integration.common_utils.managers.cc_pair import CCPairManager
-from tests.integration.common_utils.managers.cc_pair import TestCCPair
 from tests.integration.common_utils.managers.document import DocumentManager
-from tests.integration.common_utils.managers.user import TestUser
 from tests.integration.common_utils.managers.user import UserManager
+from tests.integration.common_utils.test_models import TestAPIKey
+from tests.integration.common_utils.test_models import TestCCPair
+from tests.integration.common_utils.test_models import TestUser
 
 
 def test_send_message_simple_with_history(reset: None) -> None:
@@ -19,14 +21,14 @@ def test_send_message_simple_with_history(reset: None) -> None:
     cc_pair_1: TestCCPair = CCPairManager.create_from_scratch(
         user_performing_action=admin_user,
     )
-    admin_user = UserManager.add_api_key_to_user(
-        user=admin_user,
+    api_key: TestAPIKey = APIKeyManager.create(
+        user_performing_action=admin_user,
     )
     LLMProviderManager.create(user_performing_action=admin_user)
     cc_pair_1 = DocumentManager.seed_and_attach_docs(
         cc_pair=cc_pair_1,
         num_docs=NUM_DOCS,
-        user_with_api_key=admin_user,
+        api_key=api_key,
     )
 
     response = requests.post(

backend/tests/integration/tests/document_set/test_syncing.py

@@ -1,10 +1,12 @@
 from danswer.server.documents.models import DocumentSource
 from tests.integration.common_utils.constants import NUM_DOCS
+from tests.integration.common_utils.managers.api_key import APIKeyManager
 from tests.integration.common_utils.managers.cc_pair import CCPairManager
 from tests.integration.common_utils.managers.document import DocumentManager
 from tests.integration.common_utils.managers.document_set import DocumentSetManager
-from tests.integration.common_utils.managers.user import TestUser
 from tests.integration.common_utils.managers.user import UserManager
+from tests.integration.common_utils.test_models import TestAPIKey
+from tests.integration.common_utils.test_models import TestUser
 from tests.integration.common_utils.vespa import TestVespaClient
 
 
@@ -15,8 +17,8 @@ def test_multiple_document_sets_syncing_same_connnector(
     admin_user: TestUser = UserManager.create(name="admin_user")
 
     # add api key to user
-    admin_user = UserManager.add_api_key_to_user(
-        user=admin_user,
+    api_key: TestAPIKey = APIKeyManager.create(
+        user_performing_action=admin_user,
     )
 
     # create connector
@@ -29,7 +31,7 @@ def test_multiple_document_sets_syncing_same_connnector(
     cc_pair_1 = DocumentManager.seed_and_attach_docs(
         cc_pair=cc_pair_1,
         num_docs=NUM_DOCS,
-        user_with_api_key=admin_user,
+        api_key=api_key,
     )
 
     # Create document sets