cognite-toolkit 0.7.67__py3-none-any.whl → 0.7.68__py3-none-any.whl

cognite_toolkit/_cdf_tk/client/api/instances.py

@@ -10,6 +10,7 @@ from cognite_toolkit._cdf_tk.client.resource_classes.data_modeling import (
     InstanceResponse,
     ViewReference,
 )
+from cognite_toolkit._cdf_tk.client.resource_classes.data_modeling._instance import InstanceSlimDefinition
 from cognite_toolkit._cdf_tk.client.resource_classes.instance_api import TypedInstanceIdentifier
 
 
@@ -33,15 +34,18 @@ class InstancesAPI(CDFResourceAPI[TypedInstanceIdentifier, InstanceRequest, Inst
     def _validate_response(self, response: SuccessResponse) -> ResponseItems[TypedInstanceIdentifier]:
         return ResponseItems[TypedInstanceIdentifier].model_validate_json(response.body)
 
-    def create(self, items: Sequence[InstanceRequest]) -> list[InstanceResponse]:
+    def create(self, items: Sequence[InstanceRequest]) -> list[InstanceSlimDefinition]:
         """Create instances in CDF.
 
         Args:
             items: List of InstanceRequest objects to create.
         Returns:
-            List of created InstanceResponse objects.
+            List of created InstanceSlimDefinition objects.
         """
-        return self._request_item_response(items, "upsert")
+        response_items: list[InstanceSlimDefinition] = []
+        for response in self._chunk_requests(items, "upsert", self._serialize_items):
+            response_items.extend(PagedResponse[InstanceSlimDefinition].model_validate_json(response.body).items)
+        return response_items
 
     def retrieve(
         self, items: Sequence[TypedInstanceIdentifier], source: ViewReference | None = None

cognite_toolkit/_cdf_tk/client/api/security_categories.py

@@ -4,14 +4,14 @@ from typing import Literal
 from cognite_toolkit._cdf_tk.client.cdf_client import CDFResourceAPI, PagedResponse, ResponseItems
 from cognite_toolkit._cdf_tk.client.cdf_client.api import Endpoint
 from cognite_toolkit._cdf_tk.client.http_client import HTTPClient, ItemsSuccessResponse, SuccessResponse
-from cognite_toolkit._cdf_tk.client.resource_classes.identifiers import InternalId
+from cognite_toolkit._cdf_tk.client.resource_classes.identifiers import InternalIdUnwrapped
 from cognite_toolkit._cdf_tk.client.resource_classes.securitycategory import (
     SecurityCategoryRequest,
     SecurityCategoryResponse,
 )
 
 
-class SecurityCategoriesAPI(CDFResourceAPI[InternalId, SecurityCategoryRequest, SecurityCategoryResponse]):
+class SecurityCategoriesAPI(CDFResourceAPI[InternalIdUnwrapped, SecurityCategoryRequest, SecurityCategoryResponse]):
     def __init__(self, http_client: HTTPClient) -> None:
         super().__init__(
             http_client=http_client,
@@ -27,8 +27,8 @@ class SecurityCategoriesAPI(CDFResourceAPI[InternalId, SecurityCategoryRequest,
     ) -> PagedResponse[SecurityCategoryResponse]:
         return PagedResponse[SecurityCategoryResponse].model_validate_json(response.body)
 
-    def _reference_response(self, response: SuccessResponse) -> ResponseItems[InternalId]:
-        return ResponseItems[InternalId].model_validate_json(response.body)
+    def _reference_response(self, response: SuccessResponse) -> ResponseItems[InternalIdUnwrapped]:
+        return ResponseItems[InternalIdUnwrapped].model_validate_json(response.body)
 
     def create(self, items: Sequence[SecurityCategoryRequest]) -> list[SecurityCategoryResponse]:
         """Create security categories in CDF.
@@ -40,7 +40,7 @@ class SecurityCategoriesAPI(CDFResourceAPI[InternalId, SecurityCategoryRequest,
         """
         return self._request_item_response(items, "create")
 
-    def delete(self, items: Sequence[InternalId]) -> None:
+    def delete(self, items: Sequence[InternalIdUnwrapped]) -> None:
         """Delete security categories from CDF.
 
         Args:

cognite_toolkit/_cdf_tk/client/api/transformations.py

@@ -1,10 +1,9 @@
 from collections.abc import Iterable, Sequence
-from typing import Any, Literal
+from typing import Literal
 
 from cognite_toolkit._cdf_tk.client.cdf_client import CDFResourceAPI, PagedResponse, ResponseItems
 from cognite_toolkit._cdf_tk.client.cdf_client.api import Endpoint
 from cognite_toolkit._cdf_tk.client.http_client import HTTPClient, ItemsSuccessResponse, SuccessResponse
-from cognite_toolkit._cdf_tk.client.request_classes.filters import ClassicFilter
 from cognite_toolkit._cdf_tk.client.resource_classes.identifiers import InternalOrExternalId
 from cognite_toolkit._cdf_tk.client.resource_classes.transformation import TransformationRequest, TransformationResponse
 
@@ -83,55 +82,39 @@ class TransformationsAPI(CDFResourceAPI[InternalOrExternalId, TransformationRequ
 
     def paginate(
         self,
-        filter: ClassicFilter | None = None,
-        is_public: bool | None = None,
         limit: int = 100,
         cursor: str | None = None,
     ) -> PagedResponse[TransformationResponse]:
         """Iterate over all transformations in CDF.
 
         Args:
-            filter: Filter by data set IDs.
-            is_public: Filter by public status.
             limit: Maximum number of items to return.
             cursor: Cursor for pagination.
 
         Returns:
             PagedResponse of TransformationResponse objects.
         """
-        filter_: dict[str, Any] = filter.dump() if filter else {}
-        if is_public is not None:
-            filter_["isPublic"] = is_public
-
         return self._paginate(
             cursor=cursor,
             limit=limit,
-            body={"filter": filter_ or None},
+            body={"filter": {}},
         )
 
     def iterate(
         self,
-        filter: ClassicFilter | None = None,
-        is_public: bool | None = None,
         limit: int = 100,
     ) -> Iterable[list[TransformationResponse]]:
         """Iterate over all transformations in CDF.
 
         Args:
-            filter: Filter by data set IDs.
-            is_public: Filter by public status.
             limit: Maximum number of items to return per page.
 
         Returns:
             Iterable of lists of TransformationResponse objects.
         """
-        filter_: dict[str, Any] = filter.dump() if filter else {}
-        if is_public is not None:
-            filter_["isPublic"] = is_public
-
         return self._iterate(
             limit=limit,
-            body={"filter": filter_ or None},
+            body={"filter": {}},
         )
 
     def list(
@@ -143,4 +126,4 @@ class TransformationsAPI(CDFResourceAPI[InternalOrExternalId, TransformationRequ
         Returns:
             List of TransformationResponse objects.
         """
-        return self._list(limit=limit)
+        return self._list(limit=limit, body={"filter": {}})

cognite_toolkit/_cdf_tk/client/api/workflow_versions.py

@@ -65,8 +65,8 @@ class WorkflowVersionsAPI(CDFResourceAPI[WorkflowVersionId, WorkflowVersionReque
         result: list[WorkflowVersionResponse] = []
         for item in items:
             endpoint = f"/workflows/{item.workflow_external_id}/versions/{item.version}"
-            retrieved = self._request_item_response([item], "retrieve", endpoint=endpoint)
-            result.extend(retrieved)
+            for response in self._chunk_requests([item], "retrieve", self._serialize_items, endpoint_path=endpoint):
+                result.append(WorkflowVersionResponse.model_validate_json(response.body))
         return result
 
     def delete(self, items: Sequence[WorkflowVersionId]) -> None:

cognite_toolkit/_cdf_tk/client/api/workflows.py

@@ -65,8 +65,8 @@ class WorkflowsAPI(CDFResourceAPI[ExternalId, WorkflowRequest, WorkflowResponse]
         result: list[WorkflowResponse] = []
         for item in items:
             endpoint = f"/workflows/{item.external_id}"
-            retrieved = self._request_item_response([item], "retrieve", endpoint=endpoint)
-            result.extend(retrieved)
+            for response in self._chunk_requests([item], "retrieve", self._serialize_items, endpoint_path=endpoint):
+                result.append(WorkflowResponse.model_validate_json(response.body))
         return result
 
     def delete(self, items: Sequence[ExternalId]) -> None:

cognite_toolkit/_cdf_tk/client/resource_classes/data_modeling/_instance.py

@@ -29,7 +29,7 @@ class InstanceSource(BaseModelObject):
 
 class InstanceRequestDefinition(InstanceDefinition, RequestResource, ABC):
     existing_version: int | None = None
-    sources: list[InstanceSource]
+    sources: list[InstanceSource] | None = None
 
 
 class InstanceResponseDefinition(InstanceDefinition, ResponseResource, Generic[T_RequestResource], ABC):
@@ -144,6 +144,24 @@ class EdgeResponse(InstanceResponseDefinition[EdgeRequest]):
         return EdgeRequest.model_validate(dumped, extra="ignore")
 
 
+class InstanceSlimDefinition(BaseModelObject):
+    """Slim version of instance definition for listing instances."""
+
+    instance_type: Literal["node", "edge"]
+    version: int
+    was_modified: bool
+    space: str
+    external_id: str
+    created_time: int
+    last_updated_time: int
+
+    def as_id(self) -> TypedNodeIdentifier | TypedEdgeIdentifier:
+        if self.instance_type == "node":
+            return TypedNodeIdentifier(space=self.space, external_id=self.external_id)
+        else:
+            return TypedEdgeIdentifier(space=self.space, external_id=self.external_id)
+
+
 InstanceRequest: TypeAlias = Annotated[
     NodeRequest | EdgeRequest,
     Field(discriminator="instance_type"),

cognite_toolkit/_cdf_tk/client/resource_classes/extraction_pipeline.py

@@ -1,4 +1,4 @@
-from typing import ClassVar
+from typing import Any, ClassVar, Literal
 
 from cognite_toolkit._cdf_tk.client._resource_base import (
     BaseModelObject,
@@ -49,6 +49,12 @@ class ExtractionPipelineRequest(ExtractionPipeline, UpdatableRequestResource):
         {"documentation", "source", "notification_config", "schedule", "description"}
     )
 
+    def as_update(self, mode: Literal["patch", "replace"]) -> dict[str, Any]:
+        update = super().as_update(mode)
+        # createdBy cannot be updated.
+        update["update"].pop("createdBy", None)
+        return update
+
 
 class ExtractionPipelineResponse(ExtractionPipeline, ResponseResource[ExtractionPipelineRequest]):
     id: int

cognite_toolkit/_cdf_tk/client/resource_classes/hosted_extractor_job.py

@@ -1,6 +1,6 @@
-from typing import Annotated, Literal
+from typing import Annotated, Any, Literal
 
-from pydantic import Field, JsonValue
+from pydantic import Field, JsonValue, field_validator
 
 from cognite_toolkit._cdf_tk.client._resource_base import (
     BaseModelObject,
@@ -116,13 +116,27 @@ class HostedExtractorJob(BaseModelObject):
     destination_id: str
     source_id: str
     format: JobFormat
-    config: MQTTConfig | KafkaConfig | RestConfig
+    config: MQTTConfig | KafkaConfig | RestConfig | None = None
 
     def as_id(self) -> ExternalId:
         return ExternalId(external_id=self.external_id)
 
-
-class HostedExtractorJobRequest(HostedExtractorJob, UpdatableRequestResource): ...
+    @field_validator("config", mode="before")
+    @classmethod
+    def empty_dict_as_none(cls, value: Any) -> Any:
+        if value == {}:
+            return None
+        return value
+
+
+class HostedExtractorJobRequest(HostedExtractorJob, UpdatableRequestResource):
+    def as_update(self, mode: Literal["patch", "replace"]) -> dict[str, Any]:
+        update_item = super().as_update(mode=mode)
+        exclude_unset = mode == "patch"
+        update_item["update"]["format"] = {
+            "set": self.format.model_dump(mode="json", exclude_none=True, by_alias=True, exclude_unset=exclude_unset)
+        }
+        return update_item
 
 
 class HostedExtractorJobResponse(HostedExtractorJob, ResponseResource[HostedExtractorJobRequest]):

cognite_toolkit/_cdf_tk/client/resource_classes/hosted_extractor_mapping.py

@@ -1,4 +1,4 @@
-from typing import Annotated, Literal
+from typing import Annotated, ClassVar, Literal
 
 from pydantic import Field
 
@@ -60,6 +60,7 @@ class HostedExtractorMapping(BaseModelObject):
 
 
 class HostedExtractorMappingRequest(HostedExtractorMapping, UpdatableRequestResource):
+    non_nullable_fields: ClassVar[frozenset[str]] = frozenset({"input"})
     input: MappingInput | None = None
 
 

cognite_toolkit/_cdf_tk/client/resource_classes/hosted_extractor_source/_auth.py

@@ -26,7 +26,7 @@ class BasicAuthenticationResponse(BasicAuthentication, AuthenticationResponseDef
 
 
 class ClientCredentialAuthentication(BaseModelObject):
-    type: Literal["clientCredential"] = "clientCredential"
+    type: Literal["clientCredentials"] = "clientCredentials"
     client_id: str
     token_url: str
     scopes: str

cognite_toolkit/_cdf_tk/client/resource_classes/hosted_extractor_source/_base.py

@@ -16,6 +16,8 @@ class SourceRequestDefinition(UpdatableRequestResource):
 
     def as_update(self, mode: Literal["patch", "replace"]) -> dict[str, Any]:
         output = super().as_update(mode)
+        # Move type from update to top-level, as the API expects it there
+        output["update"].pop("type")
         output["type"] = self.type
         return output
 

cognite_toolkit/_cdf_tk/client/resource_classes/hosted_extractor_source/_eventhub.py

@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import ClassVar, Literal
 
 from cognite_toolkit._cdf_tk.client._resource_base import (
     BaseModelObject,
@@ -10,13 +10,14 @@ from ._base import SourceRequestDefinition, SourceResponseDefinition
 
 
 class EventHubSource(BaseModelObject):
-    type: Literal["eventHub"] = "eventHub"
+    type: Literal["eventhub"] = "eventhub"
     host: str
     event_hub_name: str
     consumer_group: str | None = None
 
 
 class EventHubSourceRequest(EventHubSource, SourceRequestDefinition):
+    non_nullable_fields: ClassVar[frozenset[str]] = frozenset({"authentication"})
     authentication: BasicAuthenticationRequest | None = None
 
 

cognite_toolkit/_cdf_tk/client/resource_classes/hosted_extractor_source/_kafka.py

@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import ClassVar, Literal
 
 from pydantic import Field
 
@@ -31,6 +31,7 @@ class KafkaSource(BaseModelObject):
 
 
 class KafkaSourceRequest(KafkaSource, SourceRequestDefinition):
+    non_nullable_fields: ClassVar[frozenset[str]] = frozenset({"use_tls"})
     authentication: (
         BasicAuthenticationRequest | ClientCredentialAuthenticationRequest | ScramShaAuthenticationRequest | None
     ) = Field(None, discriminator="type")

cognite_toolkit/_cdf_tk/client/resource_classes/hosted_extractor_source/_mqtt.py

@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import ClassVar, Literal
 
 from cognite_toolkit._cdf_tk.client._resource_base import (
     BaseModelObject,
@@ -18,6 +18,7 @@ class MQTTSource(BaseModelObject):
 
 
 class MQTTSourceRequest(MQTTSource, SourceRequestDefinition):
+    non_nullable_fields: ClassVar[frozenset[str]] = frozenset({"use_tls"})
     authentication: BasicAuthenticationRequest | None = None
     ca_certificate: CACertificateRequest | None = None
     auth_certificate: AuthCertificateRequest | None = None

cognite_toolkit/_cdf_tk/client/resource_classes/hosted_extractor_source/_rest.py

@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import ClassVar, Literal
 
 from pydantic import Field
 
@@ -26,6 +26,7 @@ class RESTSource(BaseModelObject):
 
 
 class RESTSourceRequest(RESTSource, SourceRequestDefinition):
+    non_nullable_fields: ClassVar[frozenset[str]] = frozenset({"scheme", "port"})
     scheme: Literal["https", "http"] | None = None
     authentication: (
         BasicAuthenticationRequest | HTTPBasicAuthenticationRequest | ClientCredentialAuthenticationRequest | None

cognite_toolkit/_cdf_tk/client/resource_classes/identifiers.py

@@ -1,7 +1,7 @@
 from collections.abc import Sequence
-from typing import Annotated, Literal
+from typing import Annotated, Any, Literal
 
-from pydantic import Field
+from pydantic import Field, model_serializer, model_validator
 
 from cognite_toolkit._cdf_tk.client._resource_base import Identifier
 
@@ -10,6 +10,24 @@ class InternalOrExternalIdDefinition(Identifier):
     type: str
 
 
+class InternalIdUnwrapped(Identifier):
+    id: int
+
+    def __str__(self) -> str:
+        return f"id={self.id}"
+
+    @model_serializer
+    def serialize(self) -> int:
+        return self.id
+
+    @model_validator(mode="before")
+    @classmethod
+    def deserialize(cls, value: Any) -> Any:
+        if isinstance(value, int):
+            return {"id": value}
+        return value
+
+
 class InternalId(InternalOrExternalIdDefinition):
     type: Literal["id"] = Field("id", exclude=True)
     id: int
@@ -21,6 +39,9 @@ class InternalId(InternalOrExternalIdDefinition):
     def __str__(self) -> str:
         return f"id={self.id}"
 
+    def as_unwrapped(self) -> InternalIdUnwrapped:
+        return InternalIdUnwrapped(id=self.id)
+
 
 class ExternalId(InternalOrExternalIdDefinition):
     type: Literal["externalId"] = Field("externalId", exclude=True)

cognite_toolkit/_cdf_tk/client/resource_classes/securitycategory.py

@@ -1,20 +1,26 @@
+from pydantic import Field
+
 from cognite_toolkit._cdf_tk.client._resource_base import (
     BaseModelObject,
     RequestResource,
     ResponseResource,
 )
 
-from .identifiers import NameId
+from .identifiers import InternalIdUnwrapped
 
 
 class SecurityCategory(BaseModelObject):
     name: str
 
-    def as_id(self) -> NameId:
-        return NameId(name=self.name)
 
+class SecurityCategoryRequest(SecurityCategory, RequestResource):
+    # This is not part of the request payload.
+    id: int | None = Field(None, exclude=True)
 
-class SecurityCategoryRequest(SecurityCategory, RequestResource): ...
+    def as_id(self) -> InternalIdUnwrapped:
+        if self.id is None:
+            raise ValueError("Cannot create identifier from SecurityCategoryRequest without id set")
+        return InternalIdUnwrapped(id=self.id)
 
 
 class SecurityCategoryResponse(SecurityCategory, ResponseResource[SecurityCategoryRequest]):
@@ -22,3 +28,6 @@ class SecurityCategoryResponse(SecurityCategory, ResponseResource[SecurityCatego
 
     def as_request_resource(self) -> SecurityCategoryRequest:
         return SecurityCategoryRequest.model_validate(self.dump(), extra="ignore")
+
+    def as_id(self) -> InternalIdUnwrapped:
+        return InternalIdUnwrapped(id=self.id)

cognite_toolkit/_cdf_tk/client/resource_classes/sequence.py

@@ -1,7 +1,9 @@
-from typing import ClassVar, Literal
+from collections import defaultdict
+from typing import Any, ClassVar, Literal
 
 from cognite_toolkit._cdf_tk.client._resource_base import (
     BaseModelObject,
+    Identifier,
     ResponseResource,
     UpdatableRequestResource,
 )
@@ -17,6 +19,19 @@ class SequenceColumn(BaseModelObject):
     value_type: Literal["STRING", "DOUBLE", "LONG"] | None = None
 
 
+class SequenceColumnRequest(SequenceColumn, UpdatableRequestResource):
+    def as_id(self) -> Identifier:
+        return ExternalId(external_id=self.external_id)
+
+
+class SequenceColumnResponse(SequenceColumn, ResponseResource[SequenceColumnRequest]):
+    created_time: int
+    last_updated_time: int
+
+    def as_request_resource(self) -> SequenceColumnRequest:
+        return SequenceColumnRequest.model_validate(self.dump(), extra="ignore")
+
+
 class Sequence(BaseModelObject):
     external_id: str | None = None
     name: str | None = None
@@ -24,7 +39,6 @@ class Sequence(BaseModelObject):
     asset_id: int | None = None
     data_set_id: int | None = None
     metadata: dict[str, str] | None = None
-    columns: list[SequenceColumn]
 
     def as_id(self) -> ExternalId:
         if self.external_id is None:
@@ -34,12 +48,40 @@ class Sequence(BaseModelObject):
 
 class SequenceRequest(Sequence, UpdatableRequestResource):
     container_fields: ClassVar[frozenset[str]] = frozenset({"metadata", "columns"})
+    columns: list[SequenceColumnRequest]
+
+    def as_update(
+        self, mode: Literal["patch", "replace"], last_columns: list[SequenceColumn] | None = None
+    ) -> dict[str, Any]:
+        output = super().as_update(mode)
+        output["update"].pop("columns", None)
+        exclude_unset = mode == "patch"
+        existing_ids = {col.external_id for col in last_columns} if last_columns is not None else set()
+        columns: dict[str, list[Any]] = defaultdict(list)
+        for col in self.columns:
+            col_dumped = col.model_dump(mode="json", by_alias=True, exclude_none=True)
+            if col_dumped == {"externalId": col.external_id}:
+                continue  # No changes to this column
+
+            if col.external_id in existing_ids:
+                columns["modify"].append(col.as_update(mode=mode))
+            else:
+                columns["add"].append(col.model_dump(mode="json", by_alias=True, exclude_unset=exclude_unset))
+        if mode == "replace":
+            current_columns_ids = {col.external_id for col in self.columns}
+            for last_col in last_columns or []:
+                if last_col.external_id not in current_columns_ids:
+                    columns["remove"].append({"externalId": last_col.external_id})
+        if columns:
+            output["update"]["columns"] = dict(columns)
+        return output
 
 
 class SequenceResponse(Sequence, ResponseResource[SequenceRequest]):
     id: int
     created_time: int
     last_updated_time: int
+    columns: list[SequenceColumnResponse]
 
     def as_request_resource(self) -> SequenceRequest:
         return SequenceRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/resource_classes/transformation.py

@@ -68,7 +68,7 @@ class ViewDataSource(DestinationDefinition):
 
 
 class RawDataSource(DestinationDefinition):
-    type: Literal["raw_tables"] = "raw_tables"
+    type: Literal["raw"] = "raw"
     database: str
     table: str
 
@@ -108,7 +108,7 @@ class Transformation(BaseModelObject):
 
 class TransformationRequest(Transformation, UpdatableRequestResource):
     container_fields: ClassVar[frozenset[str]] = frozenset({"tags"})
-    non_nullable_fields: ClassVar[frozenset[str]] = frozenset({"is_public", "query", "destination"})
+    non_nullable_fields: ClassVar[frozenset[str]] = frozenset({"is_public", "query", "destination", "conflict_mode"})
     query: str | None = None
     conflict_mode: str | None = None
     destination: Destination | None = None
@@ -126,7 +126,8 @@ class TransformationResponse(Transformation, ResponseResource[TransformationRequ
     conflict_mode: str
     destination: Destination
     blocked: BlockedInfo | None = None
-    owner: str
+    # In the API spec, owner is documented as a string, but the response is actually {"user": '<uuid>'}
+    owner: str | dict[str, str]
     owner_is_current_user: bool
     has_source_oidc_credentials: bool
     has_destination_oidc_credentials: bool

cognite_toolkit/_cdf_tk/client/testing.py

@@ -13,6 +13,7 @@ from cognite.client.testing import CogniteClientMock
 from rich.console import Console
 
 from cognite_toolkit._cdf_tk.client._toolkit_client import ToolkitClient
+from cognite_toolkit._cdf_tk.client.api.hosted_extractors import HostedExtractorsAPI
 from cognite_toolkit._cdf_tk.client.api.legacy.canvas import CanvasAPI, IndustrialCanvasAPI
 from cognite_toolkit._cdf_tk.client.api.legacy.charts import ChartsAPI
 from cognite_toolkit._cdf_tk.client.api.legacy.dml import DMLAPI
@@ -38,6 +39,10 @@ from .api.datasets import DataSetsAPI
 from .api.events import EventsAPI
 from .api.extraction_pipelines import ExtractionPipelinesAPI
 from .api.filemetadata import FileMetadataAPI
+from .api.hosted_extractor_destinations import HostedExtractorDestinationsAPI
+from .api.hosted_extractor_jobs import HostedExtractorJobsAPI
+from .api.hosted_extractor_mappings import HostedExtractorMappingsAPI
+from .api.hosted_extractor_sources import HostedExtractorSourcesAPI
 from .api.infield import InfieldAPI, InFieldCDMConfigAPI, InfieldConfigAPI
 from .api.labels import LabelsAPI
 from .api.lookup import (
@@ -160,6 +165,11 @@ class ToolkitClientMock(CogniteClientMock):
         self.tool.simulators.models = MagicMock(spec_set=SimulatorModelsAPI)
         self.tool.datasets = MagicMock(spec_set=DataSetsAPI)
         self.tool.extraction_pipelines = MagicMock(spec_set=ExtractionPipelinesAPI)
+        self.tool.hosted_extractors = MagicMock(spec=HostedExtractorsAPI)
+        self.tool.hosted_extractors.sources = MagicMock(spec_set=HostedExtractorSourcesAPI)
+        self.tool.hosted_extractors.jobs = MagicMock(spec_set=HostedExtractorJobsAPI)
+        self.tool.hosted_extractors.destinations = MagicMock(spec_set=HostedExtractorDestinationsAPI)
+        self.tool.hosted_extractors.mappings = MagicMock(spec_set=HostedExtractorMappingsAPI)
         self.tool.labels = MagicMock(spec_set=LabelsAPI)
         self.tool.security_categories = MagicMock(spec_set=SecurityCategoriesAPI)
         self.tool.sequences = MagicMock(spec_set=SequencesAPI)