fix test

cognitedata · Dec 3, 2024 · c56d5de · c56d5de
1 parent ec09f63
commit c56d5de
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 6 deletions.
diff --git a/cognite/client/_api/ai/__init__.py b/cognite/client/_api/ai/__init__.py
@@ -2,7 +2,7 @@
 
 from typing import TYPE_CHECKING
 
-from cognite.client._api.ai.tools import ToolsAPI
+from cognite.client._api.ai.tools import AIToolsAPI
 from cognite.client._api_client import APIClient
 
 if TYPE_CHECKING:
@@ -12,4 +12,4 @@
 class AIAPI(APIClient):
     def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: CogniteClient) -> None:
         super().__init__(config, api_version, cognite_client)
-        self.tools = ToolsAPI(config, api_version, cognite_client)
+        self.tools = AIToolsAPI(config, api_version, cognite_client)
diff --git a/cognite/client/_api/ai/tools/__init__.py b/cognite/client/_api/ai/tools/__init__.py
@@ -2,14 +2,14 @@
 
 from typing import TYPE_CHECKING
 
-from cognite.client._api.ai.tools.documents import DocumentsAPI
+from cognite.client._api.ai.tools.documents import AIDocumentsAPI
 from cognite.client._api_client import APIClient
 
 if TYPE_CHECKING:
     from cognite.client import ClientConfig, CogniteClient
 
 
-class ToolsAPI(APIClient):
+class AIToolsAPI(APIClient):
     def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: CogniteClient) -> None:
         super().__init__(config, api_version, cognite_client)
-        self.documents = DocumentsAPI(config, api_version, cognite_client)
+        self.documents = AIDocumentsAPI(config, api_version, cognite_client)
diff --git a/cognite/client/_api/ai/tools/documents.py b/cognite/client/_api/ai/tools/documents.py
@@ -8,7 +8,7 @@
 from cognite.client.utils.useful_types import SequenceNotStr
 
 
-class DocumentsAPI(APIClient):
+class AIDocumentsAPI(APIClient):
     _RESOURCE_PATH = "/ai/tools/documents"
 
     def summarize(

diff --git a/cognite/client/testing.py b/cognite/client/testing.py
@@ -6,6 +6,9 @@
 from unittest.mock import MagicMock
 
 from cognite.client import CogniteClient
+from cognite.client._api.ai import AIAPI
+from cognite.client._api.ai.tools import AIToolsAPI
+from cognite.client._api.ai.tools.documents import AIDocumentsAPI
 from cognite.client._api.annotations import AnnotationsAPI
 from cognite.client._api.assets import AssetsAPI
 from cognite.client._api.data_modeling import DataModelingAPI
@@ -95,6 +98,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         #   - Add spacing above and below
         #   - Use `spec=MyAPI` only for "top level"
         #   - Use `spec_set=MyNestedAPI` for all nested APIs
+        self.ai = MagicMock(spec=AIAPI)
+        self.ai.tools = MagicMock(spec=AIToolsAPI)
+        self.ai.tools.documents = MagicMock(spec_set=AIDocumentsAPI)
         self.annotations = MagicMock(spec_set=AnnotationsAPI)
         self.assets = MagicMock(spec_set=AssetsAPI)