From f016a1c55081743452bf3c70dca30e08f846ab80 Mon Sep 17 00:00:00 2001 From: abdullah-cognite <100700755+abdullah-cognite@users.noreply.github.com> Date: Wed, 18 Dec 2024 12:36:53 +0000 Subject: [PATCH] first draft --- cognite/client/_api/simulators/__init__.py | 2 + .../_api/simulators/simulator_models.py | 169 ++++++++ cognite/client/_cognite_client.py | 2 + .../data_classes/simulators/__init__.py | 14 + .../client/data_classes/simulators/filters.py | 23 ++ .../client/data_classes/simulators/models.py | 365 ++++++++++++++++++ cognite/client/testing.py | 2 + .../test_api/test_simulators/conftest.py | 99 +++++ .../test_api/test_simulators/seed/data.py | 31 ++ .../test_api/test_simulators/test_models.py | 37 ++ .../test_simulators/test_simulators.py | 18 - 11 files changed, 744 insertions(+), 18 deletions(-) create mode 100644 cognite/client/_api/simulators/simulator_models.py create mode 100644 cognite/client/data_classes/simulators/filters.py create mode 100644 cognite/client/data_classes/simulators/models.py create mode 100644 tests/tests_integration/test_api/test_simulators/conftest.py create mode 100644 tests/tests_integration/test_api/test_simulators/test_models.py diff --git a/cognite/client/_api/simulators/__init__.py b/cognite/client/_api/simulators/__init__.py index aa310331c..4eb8f65e4 100644 --- a/cognite/client/_api/simulators/__init__.py +++ b/cognite/client/_api/simulators/__init__.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING +from cognite.client._api.simulators.simulator_models import SimulatorModelsAPI from cognite.client._api_client import APIClient from cognite.client._constants import DEFAULT_LIMIT_READ from cognite.client.data_classes.simulators.simulators import Simulator, SimulatorList @@ -20,6 +21,7 @@ def __init__(self, config: ClientConfig, api_version: str | None, cognite_client self._warning = FeaturePreviewWarning( api_maturity="General Availability", sdk_maturity="alpha", feature_name="Simulators" ) + self.models = SimulatorModelsAPI(config, api_version, cognite_client) def list(self, limit: int = DEFAULT_LIMIT_READ) -> SimulatorList: """`Filter simulators `_ diff --git a/cognite/client/_api/simulators/simulator_models.py b/cognite/client/_api/simulators/simulator_models.py new file mode 100644 index 000000000..593351d79 --- /dev/null +++ b/cognite/client/_api/simulators/simulator_models.py @@ -0,0 +1,169 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, overload + +from cognite.client._api_client import APIClient +from cognite.client._constants import DEFAULT_LIMIT_READ +from cognite.client.data_classes._base import CogniteFilter +from cognite.client.data_classes.simulators.filters import SimulatorModelRevisionsFilter, SimulatorModelsFilter +from cognite.client.data_classes.simulators.models import ( + CreatedTimeSort, + SimulatorModel, + SimulatorModelList, + SimulatorModelRevision, + SimulatorModelRevisionList, +) +from cognite.client.utils._experimental import FeaturePreviewWarning +from cognite.client.utils._identifier import IdentifierSequence +from cognite.client.utils.useful_types import SequenceNotStr + +if TYPE_CHECKING: + from cognite.client import ClientConfig, CogniteClient + + +class SimulatorModelRevisionsAPI(APIClient): + _RESOURCE_PATH = "/simulators/models/revisions" + + def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: CogniteClient) -> None: + super().__init__(config, api_version, cognite_client) + self._warning = FeaturePreviewWarning( + api_maturity="General Availability", sdk_maturity="alpha", feature_name="Simulators" + ) + + def list( + self, + limit: int = DEFAULT_LIMIT_READ, + sort: CreatedTimeSort | None = None, + filter: SimulatorModelRevisionsFilter | dict[str, Any] | None = None, + ) -> SimulatorModelRevisionList: + """`Filter simulator model revisions `_ + Retrieves a list of simulator model revisions that match the given criteria + Args: + limit (int): Maximum number of results to return. Defaults to 25. Set to -1, float(“inf”) or None to return all items. sort (CreatedTimeSort | None): Sort order for the results. + sort (CreatedTimeSort | None): No description. + filter (SimulatorModelRevisionsFilter | dict[str, Any] | None): Filter to apply. + Returns: + SimulatorModelRevisionList: List of simulator model revisions + Examples: + List simulator model revisions: + >>> from cognite.client import CogniteClient + >>> client = CogniteClient() + >>> res = client.simulators.models.revisions.list() + """ + self._warning.warn() + return self._list( + method="POST", + limit=limit, + resource_cls=SimulatorModelRevision, + list_cls=SimulatorModelRevisionList, + sort=[CreatedTimeSort.load(sort).dump()] if sort else None, + filter=filter.dump(camel_case=True) if isinstance(filter, CogniteFilter) else None, + ) + + @overload + def retrieve( + self, + external_id: str, + id: int | None = None, + ) -> SimulatorModelRevision | None: ... + + @overload + def retrieve(self, external_id: SequenceNotStr[str], id: int | None = None) -> SimulatorModelRevisionList: ... + + def retrieve( + self, external_id: str | SequenceNotStr[str], id: int | None = None + ) -> SimulatorModelRevision | SimulatorModelRevisionList | None: + """`Retrieve simulator model revision `_ + Retrieve a simulator model revision by ID or external ID + Args: + external_id (str | SequenceNotStr[str]): The external id of the simulator model revision. + id (int | None): The id of the simulator model revision. + Returns: + SimulatorModelRevision | SimulatorModelRevisionList | None: Requested simulator model revision + Examples: + Get simulator model revision by id: + >>> from cognite.client import CogniteClient + >>> client = CogniteClient() + >>> res = client.simulators.models.revisions.retrieve(id=123) + Get simulator model revision by external id: + >>> res = client.simulators.models.revisions.retrieve(external_id="abcdef") + """ + identifiers = IdentifierSequence.load(ids=id, external_ids=external_id).as_singleton() + return self._retrieve_multiple( + list_cls=SimulatorModelRevisionList, + resource_cls=SimulatorModelRevision, + identifiers=identifiers, + ) + + +class SimulatorModelsAPI(APIClient): + _RESOURCE_PATH = "/simulators/models" + + def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: CogniteClient) -> None: + super().__init__(config, api_version, cognite_client) + self.revisions = SimulatorModelRevisionsAPI(config, api_version, cognite_client) + self._warning = FeaturePreviewWarning( + api_maturity="General Availability", sdk_maturity="alpha", feature_name="Simulators" + ) + self._RETRIEVE_LIMIT = 1 + + def list( + self, + limit: int = DEFAULT_LIMIT_READ, + filter: SimulatorModelsFilter | dict[str, Any] | None = None, + sort: CreatedTimeSort | None = None, + ) -> SimulatorModelList: + """`Filter simulator models `_ + Retrieves a list of simulator models that match the given criteria + Args: + limit (int): Maximum number of results to return. Defaults to 25. Set to -1, float(“inf”) or None to return all items. + filter (SimulatorModelsFilter | dict[str, Any] | None): Filter to apply. + sort (CreatedTimeSort | None): The criteria to sort by. + Returns: + SimulatorModelList: List of simulator models + + Examples: + List simulator models: + >>> from cognite.client import CogniteClient + >>> client = CogniteClient() + >>> res = client.simulators.models.list() + + Specify filter and sort order: + >>> res = client.simulators.models.list( + ... filter={"name": "my_simulator_model"}, + ... sort=("created_time") + ... ) + + """ + self._warning.warn() + return self._list( + method="POST", + limit=limit, + resource_cls=SimulatorModel, + list_cls=SimulatorModelList, + sort=[CreatedTimeSort.load(sort).dump()] if sort else None, + filter=filter.dump(camel_case=True) if isinstance(filter, CogniteFilter) else None, + ) + + def retrieve(self, id: int | None = None, external_id: str | None = None) -> SimulatorModel | None: + """`Retrieve simulator model `_ + Retrieve a simulator model by ID or external ID + Args: + id (int | None): The id of the simulator model. + external_id (str | None): The external id of the simulator model. + Returns: + SimulatorModel | None: Requested simulator model + Examples: + Retrieve simulator model by id: + >>> from cognite.client import CogniteClient + >>> client = CogniteClient() + >>> res = client.simulators.models.retrieve(id=1) + Retrieve simulator model by external id: + >>> res = client.simulators.models.retrieve(external_id="foo") + """ + identifiers = IdentifierSequence.load(ids=id, external_ids=external_id).as_singleton() + return self._retrieve_multiple( + list_cls=SimulatorModelList, + resource_cls=SimulatorModel, + identifiers=identifiers, + ) diff --git a/cognite/client/_cognite_client.py b/cognite/client/_cognite_client.py index 9c9c0a410..362c3651d 100644 --- a/cognite/client/_cognite_client.py +++ b/cognite/client/_cognite_client.py @@ -24,6 +24,7 @@ from cognite.client._api.relationships import RelationshipsAPI from cognite.client._api.sequences import SequencesAPI from cognite.client._api.simulators import SimulatorsAPI +from cognite.client._api.simulators.simulator_models import SimulatorModelsAPI from cognite.client._api.templates import TemplatesAPI from cognite.client._api.three_d import ThreeDAPI from cognite.client._api.time_series import TimeSeriesAPI @@ -85,6 +86,7 @@ def __init__(self, config: ClientConfig | None = None) -> None: self.workflows = WorkflowAPI(self._config, self._API_VERSION, self) self.units = UnitAPI(self._config, self._API_VERSION, self) self.simulators = SimulatorsAPI(self._config, self._API_VERSION, self) + self.simulators.models = SimulatorModelsAPI(self._config, self._API_VERSION, self) # APIs just using base_url: self._api_client = APIClient(self._config, api_version=None, cognite_client=self) diff --git a/cognite/client/data_classes/simulators/__init__.py b/cognite/client/data_classes/simulators/__init__.py index 61743d4eb..1b3f26af6 100644 --- a/cognite/client/data_classes/simulators/__init__.py +++ b/cognite/client/data_classes/simulators/__init__.py @@ -1,5 +1,13 @@ from __future__ import annotations +from cognite.client.data_classes.simulators.models import ( + CreatedTimeSort, + PropertySort, + SimulatorModel, + SimulatorModelList, + SimulatorModelRevision, + SimulatorModelRevisionList, +) from cognite.client.data_classes.simulators.simulators import ( Simulator, SimulatorList, @@ -12,8 +20,14 @@ ) __all__ = [ + "CreatedTimeSort", + "PropertySort", "Simulator", "SimulatorList", + "SimulatorModel", + "SimulatorModelList", + "SimulatorModelRevision", + "SimulatorModelRevisionList", "SimulatorStep", "SimulatorStepField", "SimulatorStepOption", diff --git a/cognite/client/data_classes/simulators/filters.py b/cognite/client/data_classes/simulators/filters.py new file mode 100644 index 000000000..8d04a2ed7 --- /dev/null +++ b/cognite/client/data_classes/simulators/filters.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from collections.abc import Sequence + +from cognite.client.data_classes._base import CogniteFilter + + +class SimulatorModelsFilter(CogniteFilter): + def __init__( + self, + simulator_external_ids: Sequence[str] | None = None, + ) -> None: + self.simulator_external_ids = simulator_external_ids + + +class SimulatorModelRevisionsFilter(CogniteFilter): + def __init__( + self, + model_external_ids: Sequence[str] | None = None, + all_versions: bool | None = None, + ) -> None: + self.model_external_ids = model_external_ids + self.all_versions = all_versions diff --git a/cognite/client/data_classes/simulators/models.py b/cognite/client/data_classes/simulators/models.py new file mode 100644 index 000000000..abd9aa5e5 --- /dev/null +++ b/cognite/client/data_classes/simulators/models.py @@ -0,0 +1,365 @@ +from __future__ import annotations + +from abc import ABC +from typing import TYPE_CHECKING, Any, Literal + +from typing_extensions import Self + +from cognite.client.data_classes._base import ( + CogniteResourceList, + CogniteSort, + ExternalIDTransformerMixin, + IdTransformerMixin, + WriteableCogniteResource, + WriteableCogniteResourceList, +) + +if TYPE_CHECKING: + from cognite.client import CogniteClient + + +class PropertySort(CogniteSort): + def dump(self, camel_case: bool = True) -> dict[str, Any]: + dumped = super().dump(camel_case=camel_case) + dumped["property"] = self.property + return dumped + + +class CreatedTimeSort(PropertySort): + def __init__( + self, + property: Literal["createdTime"] = "createdTime", + order: Literal["asc", "desc"] = "asc", + ): + super().__init__(property, order) + + +class SimulatorModelRevisionCore(WriteableCogniteResource["SimulatorModelRevisionWrite"], ABC): + def __init__( + self, + external_id: str | None = None, + simulator_external_id: str | None = None, + model_external_id: str | None = None, + data_set_id: int | None = None, + file_id: int | None = None, + created_by_user_id: str | None = None, + status: str | None = None, + created_time: int | None = None, + last_updated_time: int | None = None, + version_number: int | None = None, + log_id: int | None = None, + description: str | None = None, + status_message: str | None = None, + ) -> None: + self.external_id = external_id + self.simulator_external_id = simulator_external_id + self.model_external_id = model_external_id + self.data_set_id = data_set_id + self.file_id = file_id + self.created_by_user_id = created_by_user_id + self.status = status + self.created_time = created_time + self.last_updated_time = last_updated_time + self.version_number = version_number + self.log_id = log_id + self.description = description + self.status_message = status_message + + @classmethod + def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self: + instance = super()._load(resource, cognite_client) + return instance + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + return super().dump(camel_case=camel_case) + + +class SimulatorModelRevisionWrite(SimulatorModelRevisionCore): + def __init__( + self, + external_id: str | None = None, + model_external_id: str | None = None, + file_id: int | None = None, + description: str | None = None, + ) -> None: + super().__init__( + external_id=external_id, + model_external_id=model_external_id, + file_id=file_id, + description=description, + ) + + def as_write(self) -> SimulatorModelRevisionWrite: + """Returns a writeable version of this resource""" + return self + + @classmethod + def _load( + cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None + ) -> SimulatorModelRevisionWrite: + return cls( + external_id=resource.get("externalId"), + model_external_id=resource.get("modelExternalId"), + file_id=resource.get("fileId"), + description=resource.get("description"), + ) + + +class SimulatorModelRevision(SimulatorModelRevisionCore): + """ + Simulator model revisions track changes and updates to a simulator model over time. + Each revision ensures that modifications to models are traceable and allows users to understand the evolution of a given model. + Args: + id (int | None): No description. + external_id (str | None): External id of the simulator model revision + simulator_external_id (str | None): No description. + model_external_id (str | None): External id of the associated simulator model + data_set_id (int | None): The id of the dataset associated with the simulator model revision + file_id (int | None): The id of the file associated with the simulator model revision + created_by_user_id (str | None): The id of the user who created the simulator model revision + status (str | None): The status of the simulator model revision + created_time (int | None): The time when the simulator model revision was created + last_updated_time (int | None): The time when the simulator model revision was last updated + version_number (int | None): The version number of the simulator model revision + log_id (int | None): The id of the log associated with the simulator model revision + description (str | None): The description of the simulator model revision + status_message (str | None): The current status of the model revision + """ + + def __init__( + self, + id: int | None = None, + external_id: str | None = None, + simulator_external_id: str | None = None, + model_external_id: str | None = None, + data_set_id: int | None = None, + file_id: int | None = None, + created_by_user_id: str | None = None, + status: str | None = None, + created_time: int | None = None, + last_updated_time: int | None = None, + version_number: int | None = None, + log_id: int | None = None, + description: str | None = None, + status_message: str | None = None, + ) -> None: + super().__init__( + external_id=external_id, + simulator_external_id=simulator_external_id, + model_external_id=model_external_id, + data_set_id=data_set_id, + file_id=file_id, + created_by_user_id=created_by_user_id, + status=status, + created_time=created_time, + last_updated_time=last_updated_time, + version_number=version_number, + log_id=log_id, + description=description, + status_message=status_message, + ) + # id/created_time/last_updated_time are required when using the class to read, + # but don't make sense passing in when creating a new object. So in order to make the typing + # correct here (i.e. int and not Optional[int]), we force the type to be int rather than + # Optional[int]. + self.id: int | None = id + self.created_time: int | None = created_time + self.last_updated_time: int | None = last_updated_time + + @classmethod + def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self: + instance = super()._load(resource, cognite_client) + return instance + + def as_write(self) -> SimulatorModelRevisionWrite: + """Returns this SimulatorModelRevision in its writing version.""" + return SimulatorModelRevisionWrite( + external_id=self.external_id, + model_external_id=self.model_external_id, + file_id=self.file_id, + description=self.description, + ) + + +class SimulatorModelCore(WriteableCogniteResource["SimulatorModelWrite"], ABC): + """ + The simulator model resource represents an asset modeled in a simulator. + This asset could range from a pump or well to a complete processing facility or refinery. + The simulator model is the root of its associated revisions, routines, runs, and results. + The dataset assigned to a model is inherited by its children. Deleting a model also deletes all its children, thereby + maintaining the integrity and hierarchy of the simulation data. + Simulator model revisions track changes and updates to a simulator model over time. + Each revision ensures that modifications to models are traceable and allows users to understand the evolution of a given model. + This is the read/response format of a simulator model. + Args: + external_id (str): External id of the simulator model + simulator_external_id (str): External id of the associated simulator + data_set_id (int): The id of the dataset associated with the simulator model + name (str): The name of the simulator model + type (str): The type key of the simulator model + description (str | None): The description of the simulator model + """ + + def __init__( + self, + external_id: str, + simulator_external_id: str, + data_set_id: int, + name: str, + type: str, + description: str | None = None, + ) -> None: + self.external_id = external_id + self.simulator_external_id = simulator_external_id + self.data_set_id = data_set_id + self.name = name + self.type = type + self.description = description + + @classmethod + def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self: + instance = super()._load(resource, cognite_client) + return instance + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + return super().dump(camel_case=camel_case) + + +class SimulatorModelWrite(SimulatorModelCore): + def __init__( + self, + external_id: str, + simulator_external_id: str, + data_set_id: int, + name: str, + type: str, + description: str | None = None, + ) -> None: + super().__init__( + external_id=external_id, + simulator_external_id=simulator_external_id, + data_set_id=data_set_id, + name=name, + type=type, + description=description, + ) + + @classmethod + def _load(cls, resource: dict, cognite_client: CogniteClient | None = None) -> SimulatorModelWrite: + return cls( + external_id=resource["externalId"], + simulator_external_id=resource["simulatorExternalId"], + data_set_id=resource["dataSetId"], + name=resource["name"], + type=resource["type"], + description=resource.get("description"), + ) + + def as_write(self) -> SimulatorModelWrite: + return self + + +class SimulatorModel(SimulatorModelCore): + """ + The simulator model resource represents an asset modeled in a simulator. + This asset could range from a pump or well to a complete processing facility or refinery. + The simulator model is the root of its associated revisions, routines, runs, and results. + The dataset assigned to a model is inherited by its children. Deleting a model also deletes all its children, thereby + maintaining the integrity and hierarchy of the simulation data. + Simulator model revisions track changes and updates to a simulator model over time. + Each revision ensures that modifications to models are traceable and allows users to understand the evolution of a given model. + This is the read/response format of a simulator model. + Args: + external_id (str): External id of the simulator model + simulator_external_id (str): External id of the associated simulator + data_set_id (int): The id of the dataset associated with the simulator model + name (str): The name of the simulator model + id (int): A unique id of a simulator model + type (str): The type key of the simulator model + description (str | None): The description of the simulator model + created_time (int | None): The time when the simulator model was created + last_updated_time (int | None): The time when the simulator model was last updated + """ + + def __init__( + self, + external_id: str, + simulator_external_id: str, + data_set_id: int, + name: str, + id: int, + type: str, + description: str | None = None, + created_time: int | None = None, + last_updated_time: int | None = None, + ) -> None: + super().__init__( + external_id=external_id, + simulator_external_id=simulator_external_id, + data_set_id=data_set_id, + name=name, + type=type, + description=description, + ) + # id/created_time/last_updated_time are required when using the class to read, + # but don't make sense passing in when creating a new object. So in order to make the typing + # correct here (i.e. int and not Optional[int]), we force the type to be int rather than + # Optional[int]. + self.id: int = id # type: ignore + self.created_time: int = created_time # type: ignore + self.last_updated_time: int = last_updated_time # type: ignore + + @classmethod + def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self: + load = super()._load(resource, cognite_client) + return cls( + external_id=load.external_id, + simulator_external_id=load.simulator_external_id, + data_set_id=load.data_set_id, + name=load.name, + id=resource["id"], + type=resource["type"], + description=load.description, + created_time=resource.get("createdTime"), + last_updated_time=resource.get("lastUpdatedTime"), + ) + + def as_write(self) -> SimulatorModelWrite: + """Returns this SimulatorModel in its writing version.""" + return SimulatorModelWrite( + external_id=self.external_id, + simulator_external_id=self.simulator_external_id, + data_set_id=self.data_set_id, + name=self.name, + type=self.type, + description=self.description, + ) + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + return super().dump(camel_case) + + +class SimulatorModelWriteList(CogniteResourceList[SimulatorModelWrite], ExternalIDTransformerMixin): + _RESOURCE = SimulatorModelWrite + + +class SimulatorModelList(WriteableCogniteResourceList[SimulatorModelWrite, SimulatorModel], IdTransformerMixin): + _RESOURCE = SimulatorModel + + def as_write(self) -> SimulatorModelWriteList: + return SimulatorModelWriteList([a.as_write() for a in self.data], cognite_client=self._get_cognite_client()) + + +class SimulatorModelRevisionWriteList(CogniteResourceList[SimulatorModelRevisionWrite], ExternalIDTransformerMixin): + _RESOURCE = SimulatorModelRevisionWrite + + +class SimulatorModelRevisionList( + WriteableCogniteResourceList[SimulatorModelRevisionWrite, SimulatorModelRevision], IdTransformerMixin +): + _RESOURCE = SimulatorModelRevision + + def as_write(self) -> SimulatorModelRevisionWriteList: + return SimulatorModelRevisionWriteList( + [a.as_write() for a in self.data], cognite_client=self._get_cognite_client() + ) diff --git a/cognite/client/testing.py b/cognite/client/testing.py index 4e0712a8c..9d8e0e2bd 100644 --- a/cognite/client/testing.py +++ b/cognite/client/testing.py @@ -44,6 +44,7 @@ from cognite.client._api.relationships import RelationshipsAPI from cognite.client._api.sequences import SequencesAPI, SequencesDataAPI from cognite.client._api.simulators import SimulatorsAPI +from cognite.client._api.simulators.simulator_models import SimulatorModelsAPI from cognite.client._api.synthetic_time_series import SyntheticDatapointsAPI from cognite.client._api.templates import ( TemplateGroupsAPI, @@ -143,6 +144,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.relationships = MagicMock(spec_set=RelationshipsAPI) self.simulators = MagicMock(spec_set=SimulatorsAPI) + self.simulators.models = MagicMock(spec_set=SimulatorModelsAPI) self.sequences = MagicMock(spec=SequencesAPI) self.sequences.data = MagicMock(spec_set=SequencesDataAPI) diff --git a/tests/tests_integration/test_api/test_simulators/conftest.py b/tests/tests_integration/test_api/test_simulators/conftest.py new file mode 100644 index 000000000..df149f40d --- /dev/null +++ b/tests/tests_integration/test_api/test_simulators/conftest.py @@ -0,0 +1,99 @@ +import time + +import pytest + +from cognite.client._cognite_client import CogniteClient +from cognite.client.data_classes.simulators.filters import SimulatorModelRevisionsFilter +from cognite.client.exceptions import CogniteAPIError +from cognite.client.utils._text import random_string +from tests.tests_integration.test_api.test_simulators.seed.data import ( + resource_names, + simulator, + simulator_integration, + simulator_model, + simulator_model_revision, +) + + +@pytest.fixture(scope="class") +def seed_resource_names() -> dict[str, str]: + resource_names["simulator_model_external_id"] += random_string() + resource_names["simulator_model_revision_external_id"] += random_string() + return resource_names + + +@pytest.fixture +def seed_simulator(cognite_client: CogniteClient, seed_resource_names) -> None: + simulator_external_id = seed_resource_names["simulator_external_id"] + simulators = cognite_client.simulators.list() + simulator_exists = len(list(filter(lambda x: x.external_id == simulator_external_id, simulators))) > 0 + if not simulator_exists: + cognite_client.post( + f"/api/v1/projects/{cognite_client.config.project}/simulators", + json={"items": [simulator]}, + ) + + +@pytest.fixture +def seed_simulator_integration(cognite_client: CogniteClient, seed_simulator, seed_resource_names) -> None: + try: + simulator_integration["heartbeat"] = int(time.time() * 1000) + cognite_client.post( + f"/api/v1/projects/{cognite_client.config.project}/simulators/integrations", + json={"items": [simulator_integration]}, + ) + except CogniteAPIError: + simulator_integrations = cognite_client.simulators.integrations.list() + integration_id = next( + filter( + lambda x: x.external_id == simulator_integration["externalId"], + simulator_integrations, + ) + ).id + # update hearbeat instead + cognite_client.post( + f"/api/v1/projects/{cognite_client.config.project}/simulators/integrations/update", + json={"items": [{"id": integration_id, "update": {"heartbeat": {"set": int(time.time() * 1000)}}}]}, + ) + + +@pytest.fixture +def seed_simulator_models(cognite_client: CogniteClient, seed_simulator_integration, seed_resource_names) -> None: + model_unique_external_id = seed_resource_names["simulator_model_external_id"] + models = cognite_client.simulators.models.list() + model_exists = len(list(filter(lambda x: x.external_id == model_unique_external_id, models))) > 0 + + if not model_exists: + cognite_client.post( + f"/api/v1/projects/{cognite_client.config.project}/simulators/models", + json={ + "items": [{**simulator_model, "externalId": model_unique_external_id}] + }, # Post actual simulator models here + ) + + +@pytest.fixture +def seed_simulator_model_revisions(cognite_client: CogniteClient, seed_simulator_models, seed_file) -> None: + model_unique_external_id = resource_names["simulator_model_external_id"] + model_revision_unique_external_id = resource_names["simulator_model_revision_external_id"] + model_revisions = cognite_client.simulators.models.revisions.list( + filter=SimulatorModelRevisionsFilter(model_external_ids=[model_unique_external_id]) + ) + model_revision_not_exists = ( + len(list(filter(lambda x: x.external_id == model_revision_unique_external_id, model_revisions))) == 0 + ) + + if model_revision_not_exists: + cognite_client.post( + f"/api/v1/projects/{cognite_client.config.project}/simulators/models/revisions", + json={ + "items": [ + { + **simulator_model_revision, + "fileId": seed_file.id, + "modelExternalId": model_unique_external_id, + "externalId": model_revision_unique_external_id, + } + ] + }, + ) diff --git a/tests/tests_integration/test_api/test_simulators/seed/data.py b/tests/tests_integration/test_api/test_simulators/seed/data.py index e376b6dac..e40db96db 100644 --- a/tests/tests_integration/test_api/test_simulators/seed/data.py +++ b/tests/tests_integration/test_api/test_simulators/seed/data.py @@ -1,3 +1,5 @@ +import time + data_set_id = 97552494921583 resource_names = { @@ -199,3 +201,32 @@ }, ], } + +simulator_model = { + "externalId": resource_names["simulator_model_external_id"], + "simulatorExternalId": resource_names["simulator_external_id"], + "name": "Test Simulator Model", + "description": "Test Simulator Model Desc", + "dataSetId": resource_names["simulator_test_data_set_id"], + "type": "SteadyState", +} + +simulator_model_revision = { + "externalId": resource_names["simulator_model_revision_external_id"], + "modelExternalId": resource_names["simulator_model_external_id"], + "description": "test sim model revision description", + "fileId": 00000000000000, +} + +simulator_integration = { + "externalId": resource_names["simulator_integration_external_id"], + "simulatorExternalId": resource_names["simulator_external_id"], + "heartbeat": int(time.time() * 1000), + "dataSetId": resource_names["simulator_test_data_set_id"], + "connectorVersion": "1.0.0", + "simulatorVersion": "1.0.0", + "licenseStatus": "AVAILABLE", + "licenseLastCheckedTime": 0, + "connectorStatus": "IDLE", + "connectorStatusUpdatedTime": 0, +} diff --git a/tests/tests_integration/test_api/test_simulators/test_models.py b/tests/tests_integration/test_api/test_simulators/test_models.py new file mode 100644 index 000000000..ab8e5f654 --- /dev/null +++ b/tests/tests_integration/test_api/test_simulators/test_models.py @@ -0,0 +1,37 @@ +import pytest + +from cognite.client._cognite_client import CogniteClient +from cognite.client.data_classes.simulators.filters import SimulatorModelRevisionsFilter, SimulatorModelsFilter + + +@pytest.mark.usefixtures( + "seed_resource_names", + "seed_simulator_model_revisions", +) +class TestSimulatorModels: + def test_list_models(self, cognite_client: CogniteClient, seed_resource_names) -> None: + models = cognite_client.simulators.models.list( + limit=5, filter=SimulatorModelsFilter(simulator_external_ids=[seed_resource_names["simulator_external_id"]]) + ) + assert len(models) > 0 + + def test_retrieve_model(self, cognite_client: CogniteClient, seed_resource_names) -> None: + model_external_id = seed_resource_names["simulator_model_external_id"] + model = cognite_client.simulators.models.retrieve(external_id=model_external_id) + assert model is not None + assert model.external_id == model_external_id + + def test_list_model_revisions(self, cognite_client: CogniteClient, seed_resource_names) -> None: + model_external_id = seed_resource_names["simulator_model_external_id"] + + revisions = cognite_client.simulators.models.revisions.list( + limit=5, + filter=SimulatorModelRevisionsFilter(model_external_ids=[model_external_id]), + ) + assert len(revisions) > 0 + + def test_retrieve_model_revision(self, cognite_client: CogniteClient, seed_resource_names) -> None: + model_revision_external_id = seed_resource_names["simulator_model_revision_external_id"] + model_revision = cognite_client.simulators.models.revisions.retrieve(external_id=model_revision_external_id) + assert model_revision is not None + assert model_revision.model_external_id == model_revision_external_id diff --git a/tests/tests_integration/test_api/test_simulators/test_simulators.py b/tests/tests_integration/test_api/test_simulators/test_simulators.py index c33b0c6f8..60789cc14 100644 --- a/tests/tests_integration/test_api/test_simulators/test_simulators.py +++ b/tests/tests_integration/test_api/test_simulators/test_simulators.py @@ -1,24 +1,6 @@ import pytest from cognite.client._cognite_client import CogniteClient -from tests.tests_integration.test_api.test_simulators.seed.data import resource_names, simulator - - -@pytest.fixture(scope="class") -def seed_resource_names() -> dict[str, str]: - return resource_names - - -@pytest.fixture -def seed_simulator(cognite_client: CogniteClient, seed_resource_names) -> None: - simulator_external_id = seed_resource_names["simulator_external_id"] - simulators = cognite_client.simulators.list() - simulator_exists = len(list(filter(lambda x: x.external_id == simulator_external_id, simulators))) > 0 - if not simulator_exists: - cognite_client.post( - f"/api/v1/projects/{cognite_client.config.project}/simulators", - json={"items": [simulator]}, - ) @pytest.mark.usefixtures("seed_resource_names", "seed_simulator")