Resolve type checking errors in tests and examples

Author: cyberkunju

examples/wait_for_knowledge_base.py

@@ -29,6 +29,10 @@ def main() -> None:
         embedding_model_uuid="your-embedding-model-uuid",  # Use your embedding model UUID
     )
 
+    if not kb_response.knowledge_base or not kb_response.knowledge_base.uuid:
+        print("Failed to create knowledge base")
+        return
+
     kb_uuid = kb_response.knowledge_base.uuid
     print(f"Created knowledge base: {kb_uuid}")
 

tests/api_resources/test_knowledge_bases.py

@@ -282,14 +282,14 @@ def test_method_wait_for_database_success(self, client: Gradient) -> None:
 
         call_count = [0]
 
-        def mock_retrieve(_uuid, **_kwargs):
+        def mock_retrieve(uuid: str, **kwargs: object) -> Mock:  # noqa: ARG001
             call_count[0] += 1
             response = Mock()
             # Simulate CREATING -> ONLINE transition
             response.database_status = "CREATING" if call_count[0] == 1 else "ONLINE"
             return response
 
-        client.knowledge_bases.retrieve = mock_retrieve
+        client.knowledge_bases.retrieve = mock_retrieve  # type: ignore[method-assign]
 
         result = client.knowledge_bases.wait_for_database(
             "test-uuid",
@@ -307,12 +307,12 @@ def test_method_wait_for_database_failed_state(self, client: Gradient) -> None:
 
         from gradient.resources.knowledge_bases import KnowledgeBaseDatabaseError
 
-        def mock_retrieve(_uuid, **_kwargs):
+        def mock_retrieve(uuid: str, **kwargs: object) -> Mock:  # noqa: ARG001
             response = Mock()
             response.database_status = "UNHEALTHY"
             return response
 
-        client.knowledge_bases.retrieve = mock_retrieve
+        client.knowledge_bases.retrieve = mock_retrieve  # type: ignore[method-assign]
 
         with pytest.raises(KnowledgeBaseDatabaseError, match="UNHEALTHY"):
             client.knowledge_bases.wait_for_database(
@@ -328,12 +328,12 @@ def test_method_wait_for_database_timeout(self, client: Gradient) -> None:
 
         from gradient.resources.knowledge_bases import KnowledgeBaseTimeoutError
 
-        def mock_retrieve(_uuid, **_kwargs):
+        def mock_retrieve(uuid: str, **kwargs: object) -> Mock:  # noqa: ARG001
             response = Mock()
             response.database_status = "CREATING"
             return response
 
-        client.knowledge_bases.retrieve = mock_retrieve
+        client.knowledge_bases.retrieve = mock_retrieve  # type: ignore[method-assign]
 
         with pytest.raises(KnowledgeBaseTimeoutError):
             client.knowledge_bases.wait_for_database(
@@ -349,12 +349,12 @@ def test_method_wait_for_database_decommissioned(self, client: Gradient) -> None
 
         from gradient.resources.knowledge_bases import KnowledgeBaseDatabaseError
 
-        def mock_retrieve(_uuid, **_kwargs):
+        def mock_retrieve(uuid: str, **kwargs: object) -> Mock:  # noqa: ARG001
             response = Mock()
             response.database_status = "DECOMMISSIONED"
             return response
 
-        client.knowledge_bases.retrieve = mock_retrieve
+        client.knowledge_bases.retrieve = mock_retrieve  # type: ignore[method-assign]
 
         with pytest.raises(KnowledgeBaseDatabaseError, match="DECOMMISSIONED"):
             client.knowledge_bases.wait_for_database(
@@ -636,14 +636,14 @@ async def test_method_wait_for_database_success(self, async_client: AsyncGradien
 
         call_count = [0]
 
-        async def mock_retrieve(_uuid, **_kwargs):
+        async def mock_retrieve(uuid: str, **kwargs: object) -> Mock:  # noqa: ARG001
             call_count[0] += 1
             response = Mock()
             # Simulate CREATING -> ONLINE transition
             response.database_status = "CREATING" if call_count[0] == 1 else "ONLINE"
             return response
 
-        async_client.knowledge_bases.retrieve = mock_retrieve
+        async_client.knowledge_bases.retrieve = mock_retrieve  # type: ignore[method-assign]
 
         result = await async_client.knowledge_bases.wait_for_database(
             "test-uuid",
@@ -661,12 +661,12 @@ async def test_method_wait_for_database_failed_state(self, async_client: AsyncGr
 
         from gradient.resources.knowledge_bases import KnowledgeBaseDatabaseError
 
-        async def mock_retrieve(_uuid, **_kwargs):
+        async def mock_retrieve(uuid: str, **kwargs: object) -> Mock:  # noqa: ARG001
             response = Mock()
             response.database_status = "UNHEALTHY"
             return response
 
-        async_client.knowledge_bases.retrieve = mock_retrieve
+        async_client.knowledge_bases.retrieve = mock_retrieve  # type: ignore[method-assign]
 
         with pytest.raises(KnowledgeBaseDatabaseError, match="UNHEALTHY"):
             await async_client.knowledge_bases.wait_for_database(
@@ -682,12 +682,12 @@ async def test_method_wait_for_database_timeout(self, async_client: AsyncGradien
 
         from gradient.resources.knowledge_bases import KnowledgeBaseTimeoutError
 
-        async def mock_retrieve(_uuid, **_kwargs):
+        async def mock_retrieve(uuid: str, **kwargs: object) -> Mock:  # noqa: ARG001
             response = Mock()
             response.database_status = "CREATING"
             return response
 
-        async_client.knowledge_bases.retrieve = mock_retrieve
+        async_client.knowledge_bases.retrieve = mock_retrieve  # type: ignore[method-assign]
 
         with pytest.raises(KnowledgeBaseTimeoutError):
             await async_client.knowledge_bases.wait_for_database(
@@ -703,12 +703,12 @@ async def test_method_wait_for_database_decommissioned(self, async_client: Async
 
         from gradient.resources.knowledge_bases import KnowledgeBaseDatabaseError
 
-        async def mock_retrieve(_uuid, **_kwargs):
+        async def mock_retrieve(uuid: str, **kwargs: object) -> Mock:  # noqa: ARG001
             response = Mock()
             response.database_status = "DECOMMISSIONED"
             return response
 
-        async_client.knowledge_bases.retrieve = mock_retrieve
+        async_client.knowledge_bases.retrieve = mock_retrieve  # type: ignore[method-assign]
 
         with pytest.raises(KnowledgeBaseDatabaseError, match="DECOMMISSIONED"):
             await async_client.knowledge_bases.wait_for_database(