From fe5df3c8eec910c9a997b7c63a0175d1f27997cf Mon Sep 17 00:00:00 2001
From: pablodanswer <pablo@danswer.ai>
Date: Thu, 19 Sep 2024 12:54:27 -0700
Subject: [PATCH] update integration test models

---
 backend/tests/integration/common_utils/llm.py        |  7 ++++---
 .../tests/integration/common_utils/test_models.py    | 12 ++++++++++++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/backend/tests/integration/common_utils/llm.py b/backend/tests/integration/common_utils/llm.py
index f74b40073c98..ce1388e1698f 100644
--- a/backend/tests/integration/common_utils/llm.py
+++ b/backend/tests/integration/common_utils/llm.py
@@ -7,6 +7,7 @@ from danswer.server.manage.llm.models import LLMProviderUpsertRequest
 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 TestLLMProvider
+from tests.integration.common_utils.test_models import TestLLMProviderResponse
 from tests.integration.common_utils.test_models import TestUser
 
 
@@ -22,7 +23,7 @@ class LLMProviderManager:
         groups: list[int] | None = None,
         is_public: bool | None = None,
         user_performing_action: TestUser | None = None,
-    ) -> TestLLMProvider:
+    ) -> TestLLMProviderResponse:
         print("Seeding LLM Providers...")
 
         llm_provider = LLMProviderUpsertRequest(
@@ -49,11 +50,11 @@ class LLMProviderManager:
         )
         llm_response.raise_for_status()
         response_data = llm_response.json()
-        result_llm = TestLLMProvider(
+        result_llm = TestLLMProviderResponse(
             id=response_data["id"],
             name=response_data["name"],
             provider=response_data["provider"],
-            api_key=response_data["api_key"],
+            api_key_set=response_data["api_key_set"],
             default_model_name=response_data["default_model_name"],
             is_public=response_data["is_public"],
             groups=response_data["groups"],
diff --git a/backend/tests/integration/common_utils/test_models.py b/backend/tests/integration/common_utils/test_models.py
index 04db0851e3d6..7d170c98af9b 100644
--- a/backend/tests/integration/common_utils/test_models.py
+++ b/backend/tests/integration/common_utils/test_models.py
@@ -91,6 +91,18 @@ class TestLLMProvider(BaseModel):
     api_version: str | None = None
 
 
+class TestLLMProviderResponse(BaseModel):
+    id: int
+    name: str
+    provider: str
+    api_key_set: bool
+    default_model_name: str
+    is_public: bool
+    groups: list[TestUserGroup]
+    api_base: str | None = None
+    api_version: str | None = None
+
+
 class TestDocumentSet(BaseModel):
     id: int
     name: str