fix(simint): retrieve(..) overloads and enable mypy checks (#2315)

Co-authored-by: erlendvollset <erlendvollset@gmail.com>

Author: polomani

.pre-commit-config.yaml

@@ -29,12 +29,15 @@ repos:
     hooks:
       - id: mypy
         name: mypy
-      # Eventually we want to run mypy on the whole test code base. But for now we only enable it for data modeling tests.
+      # Eventually we want to run mypy on the whole test code base. But for now we only enable it for data modeling, simulators and datapoints
         entry: |
           poetry run dmypy run -- cognite tests/tests_integration/test_api/test_data_modeling
           tests/tests_unit/test_data_classes/test_data_models
           tests/tests_unit/test_utils/test_time.py
           tests/tests_integration/test_api/test_datapoints.py
+          tests/tests_integration/test_api/test_simulators
+          tests/tests_unit/test_api/test_simulators
+          tests/tests_unit/test_data_classes/test_simulators.py
         files: ^.*.(py|pyi)$
         language: system
         pass_filenames: false

cognite/client/_api/simulators/__init__.py

@@ -70,11 +70,11 @@ def __call__(
             limit=limit,
         )
 
-    def list(self, limit: int = DEFAULT_LIMIT_READ) -> SimulatorList:
+    def list(self, limit: int | None = DEFAULT_LIMIT_READ) -> SimulatorList:
         """`List all simulators <https://developer.cognite.com/api#tag/Simulators/operation/filter_simulators_simulators_list_post>`_
 
         Args:
-            limit (int): Maximum number of results to return. Defaults to 25. Set to -1, float(“inf”) or None to return all items.
+            limit (int | None): Maximum number of results to return. Defaults to 25. Set to -1, float(“inf”) or None to return all items.
 
         Returns:
             SimulatorList: List of simulators

cognite/client/_api/simulators/models.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterator, Sequence
-from typing import TYPE_CHECKING, NoReturn, overload
+from typing import TYPE_CHECKING, overload
 
 from cognite.client._api.simulators.models_revisions import SimulatorModelRevisionsAPI
 from cognite.client._api_client import APIClient
@@ -37,7 +37,7 @@ def __init__(self, config: ClientConfig, api_version: str | None, cognite_client
 
     def list(
         self,
-        limit: int = DEFAULT_LIMIT_READ,
+        limit: int | None = DEFAULT_LIMIT_READ,
         simulator_external_ids: str | SequenceNotStr[str] | None = None,
         sort: PropertySort | None = None,
     ) -> SimulatorModelList:
@@ -46,7 +46,7 @@ def list(
         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.
+            limit (int | None): Maximum number of results to return. Defaults to 25. Set to -1, float(“inf”) or None to return all items.
             simulator_external_ids (str | SequenceNotStr[str] | None): Filter by simulator external id(s).
             sort (PropertySort | None): The criteria to sort by.
 
@@ -82,27 +82,20 @@ def list(
         )
 
     @overload
-    def retrieve(self, ids: None = None, external_ids: None = None) -> NoReturn: ...
+    def retrieve(self, *, ids: int) -> SimulatorModel | None: ...
 
     @overload
-    def retrieve(self, ids: int, external_ids: None = None) -> SimulatorModel | None: ...
+    def retrieve(self, *, external_ids: str) -> SimulatorModel | None: ...
 
     @overload
-    def retrieve(
-        self,
-        ids: None,
-        external_ids: str,
-    ) -> SimulatorModel | None: ...
+    def retrieve(self, *, ids: Sequence[int]) -> SimulatorModelList: ...
 
     @overload
-    def retrieve(
-        self,
-        ids: Sequence[int] | None = None,
-        external_ids: SequenceNotStr[str] | None = None,
-    ) -> SimulatorModelList | None: ...
+    def retrieve(self, *, external_ids: SequenceNotStr[str]) -> SimulatorModelList: ...
 
     def retrieve(
         self,
+        *,
         ids: int | Sequence[int] | None = None,
         external_ids: str | SequenceNotStr[str] | None = None,
     ) -> SimulatorModel | SimulatorModelList | None:

cognite/client/_api/simulators/models_revisions.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterator, Sequence
-from typing import TYPE_CHECKING, NoReturn, overload
+from typing import TYPE_CHECKING, overload
 
 from cognite.client._api_client import APIClient
 from cognite.client._constants import DEFAULT_LIMIT_READ
@@ -90,27 +90,20 @@ def list(
         )
 
     @overload
-    def retrieve(self, ids: None = None, external_ids: None = None) -> NoReturn: ...
+    def retrieve(self, *, ids: int) -> SimulatorModelRevision | None: ...
 
     @overload
-    def retrieve(self, ids: int, external_ids: None = None) -> SimulatorModelRevision | None: ...
+    def retrieve(self, *, external_ids: str) -> SimulatorModelRevision | None: ...
 
     @overload
-    def retrieve(
-        self,
-        ids: None,
-        external_ids: str,
-    ) -> SimulatorModelRevision | None: ...
+    def retrieve(self, *, ids: Sequence[int]) -> SimulatorModelRevisionList: ...
 
     @overload
-    def retrieve(
-        self,
-        ids: int | Sequence[int] | None = None,
-        external_ids: str | SequenceNotStr[str] | None = None,
-    ) -> SimulatorModelRevision | SimulatorModelRevisionList | None: ...
+    def retrieve(self, *, external_ids: SequenceNotStr[str]) -> SimulatorModelRevisionList: ...
 
     def retrieve(
         self,
+        *,
         ids: int | Sequence[int] | None = None,
         external_ids: str | SequenceNotStr[str] | None = None,
     ) -> SimulatorModelRevision | SimulatorModelRevisionList | None:

cognite/client/_api/simulators/routine_revisions.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterator, Sequence
-from typing import TYPE_CHECKING, NoReturn, overload
+from typing import TYPE_CHECKING, overload
 
 from cognite.client._api_client import APIClient
 from cognite.client.data_classes.shared import TimestampRange
@@ -127,27 +127,20 @@ def __call__(
         )
 
     @overload
-    def retrieve(self, ids: None = None, external_ids: None = None) -> NoReturn: ...
+    def retrieve(self, *, ids: int) -> SimulatorRoutineRevision | None: ...
 
     @overload
-    def retrieve(self, ids: int, external_ids: None = None) -> SimulatorRoutineRevision | None: ...
+    def retrieve(self, *, external_ids: str) -> SimulatorRoutineRevision | None: ...
 
     @overload
-    def retrieve(
-        self,
-        ids: None,
-        external_ids: str,
-    ) -> SimulatorRoutineRevision | None: ...
+    def retrieve(self, *, ids: Sequence[int]) -> SimulatorRoutineRevisionList: ...
 
     @overload
-    def retrieve(
-        self,
-        ids: Sequence[int] | None = None,
-        external_ids: SequenceNotStr[str] | None = None,
-    ) -> SimulatorRoutineRevisionList | None: ...
+    def retrieve(self, *, external_ids: SequenceNotStr[str]) -> SimulatorRoutineRevisionList: ...
 
     def retrieve(
         self,
+        *,
         ids: int | Sequence[int] | None = None,
         external_ids: str | SequenceNotStr[str] | None = None,
     ) -> SimulatorRoutineRevision | SimulatorRoutineRevisionList | None:

tests/tests_integration/test_api/test_simulators/conftest.py

@@ -4,23 +4,31 @@
 import time
 from collections.abc import Iterator
 from pathlib import Path
-from typing import Any
 
 import pytest
 
 from cognite.client._cognite_client import CogniteClient
 from cognite.client.data_classes.data_sets import DataSetWrite
 from cognite.client.data_classes.files import FileMetadata
 from cognite.client.data_classes.simulators import SimulatorModelWrite
-from cognite.client.data_classes.simulators.routine_revisions import SimulatorRoutineRevisionWrite
-from cognite.client.data_classes.simulators.routines import SimulatorRoutineWrite
+from cognite.client.data_classes.simulators.models import SimulatorModel
+from cognite.client.data_classes.simulators.routine_revisions import (
+    SimulatorRoutineRevision,
+    SimulatorRoutineRevisionWrite,
+)
+from cognite.client.data_classes.simulators.routines import (
+    SimulatorRoutine,
+    SimulatorRoutineList,
+    SimulatorRoutineWrite,
+)
+from cognite.client.utils._text import to_snake_case
 from tests.tests_integration.test_api.test_simulators.seed.data import (
     RESOURCES,
     SIMULATOR,
     SIMULATOR_INTEGRATION,
     SIMULATOR_MODEL,
     SIMULATOR_ROUTINE,
-    SIMULATOR_ROUTINE_REVISION_OBJ,
+    SIMULATOR_ROUTINE_REVISION,
     ResourceNames,
 )
 from tests.tests_integration.test_api.test_simulators.utils import update_logs
@@ -44,12 +52,14 @@ def seed_resource_names(cognite_client: CogniteClient) -> ResourceNames:
 def upload_file(cognite_client: CogniteClient, filename: str, external_id: str, data_set_id: int) -> FileMetadata:
     file = cognite_client.files.retrieve(external_id=external_id)
     if not file:
-        return cognite_client.files.upload(
-            path=SEED_DIR / filename,
+        uploaded_file = cognite_client.files.upload(
+            path=str(SEED_DIR / filename),
             external_id=external_id,
             name=filename,
             data_set_id=data_set_id,
         )
+        assert isinstance(uploaded_file, FileMetadata)
+        return uploaded_file
 
     return file
 
@@ -70,46 +80,47 @@ def seed_model_revision_file(
 
 
 @pytest.fixture(scope="session")
-def seed_simulator(cognite_client: CogniteClient, seed_resource_names: ResourceNames) -> Iterator[None]:
+def seed_simulator(cognite_client: CogniteClient, seed_resource_names: ResourceNames) -> None:
     simulator_external_id = seed_resource_names.simulator_external_id
     simulators = cognite_client.simulators.list(limit=None)
     seeded_simulator = simulators.get(external_id=simulator_external_id)
     fields_to_compare = ["fileExtensionTypes", "modelTypes", "modelDependencies", "stepFields", "unitQuantities"]
-    seeded_simulator_dump = seeded_simulator.dump() if seeded_simulator else None
 
     if not seeded_simulator:
         cognite_client.simulators._post("/simulators", json={"items": [SIMULATOR]})
     # if any field in simulator is different from the current seeded simulator, update it
-    elif any(seeded_simulator_dump.get(field) != SIMULATOR[field] for field in fields_to_compare if field in SIMULATOR):
+    elif any(getattr(seeded_simulator, to_snake_case(field)) != SIMULATOR.get(field) for field in fields_to_compare):
         simulator_update = {
             "id": seeded_simulator.id,
-            "update": {field: {"set": SIMULATOR[field]} for field in fields_to_compare},
+            "update": {field: {"set": SIMULATOR.get(field)} for field in fields_to_compare},
         }
         cognite_client.simulators._post("/simulators/update", json={"items": [simulator_update]})
 
 
 @pytest.fixture(scope="session")
 def seed_simulator_integration(
     cognite_client: CogniteClient, seed_simulator: None, seed_resource_names: ResourceNames
-) -> Iterator[None]:
+) -> None:
     log_id = None
     timestamp = int(time.time() * 1000)
     simulator_integrations = cognite_client.simulators.integrations.list(limit=None)
-    if not simulator_integrations.get(external_id=SIMULATOR_INTEGRATION["externalId"]):
-        SIMULATOR_INTEGRATION["heartbeat"] = timestamp
-        SIMULATOR_INTEGRATION["dataSetId"] = seed_resource_names.simulator_test_data_set_id
+    existing_integration = simulator_integrations.get(external_id=seed_resource_names.simulator_integration_external_id)
+    if not existing_integration:
+        new_integration = {
+            **SIMULATOR_INTEGRATION,
+            "heartbeat": timestamp,
+            "dataSetId": seed_resource_names.simulator_test_data_set_id,
+        }
         res = cognite_client.simulators._post(
             "/simulators/integrations",
-            json={"items": [SIMULATOR_INTEGRATION]},
+            json={"items": [new_integration]},
         )
         log_id = res.json()["items"][0]["logId"]
     else:
-        integration = simulator_integrations.get(external_id=SIMULATOR_INTEGRATION["externalId"])
-        if integration is not None:
-            log_id = integration.log_id
+        log_id = existing_integration.log_id
         cognite_client.simulators.integrations._post(
             "/simulators/integrations/update",
-            json={"items": [{"id": integration.id, "update": {"heartbeat": {"set": timestamp}}}]},
+            json={"items": [{"id": existing_integration.id, "update": {"heartbeat": {"set": timestamp}}}]},
         )
 
     if log_id:
@@ -123,35 +134,33 @@ def seed_simulator_integration(
 @pytest.fixture(scope="session")
 def seed_simulator_models(
     cognite_client: CogniteClient, seed_simulator_integration: None, seed_resource_names: ResourceNames
-) -> Iterator[dict[str, Any]]:
+) -> Iterator[SimulatorModel]:
     model_unique_external_id = seed_resource_names.simulator_model_external_id
     models = cognite_client.simulators.models.list(limit=None)
-    model_exists = models.get(external_id=model_unique_external_id)
-
-    if not model_exists:
-        SIMULATOR_MODEL["dataSetId"] = seed_resource_names.simulator_test_data_set_id
-        SIMULATOR_MODEL["externalId"] = model_unique_external_id
+    model = models.get(external_id=model_unique_external_id)
 
-        model = SimulatorModelWrite._load(
+    if not model:
+        new_model = SimulatorModelWrite._load(
             {
-                "externalId": SIMULATOR_MODEL["externalId"],
-                "simulatorExternalId": SIMULATOR_MODEL["simulatorExternalId"],
-                "dataSetId": SIMULATOR_MODEL["dataSetId"],
-                "name": SIMULATOR_MODEL["name"],
-                "type": SIMULATOR_MODEL["type"],
-                "description": SIMULATOR_MODEL["description"],
+                **SIMULATOR_MODEL.dump(),
+                "dataSetId": seed_resource_names.simulator_test_data_set_id,
+                "externalId": model_unique_external_id,
             }
         )
-        cognite_client.simulators.models.create(model)
-
-    yield SIMULATOR_MODEL
+        res = cognite_client.simulators.models.create(new_model)
+        yield res
+    else:
+        yield model
 
     cognite_client.simulators.models.delete(external_ids=model_unique_external_id)
 
 
 @pytest.fixture(scope="session")
 def seed_simulator_model_revisions(
-    cognite_client: CogniteClient, seed_simulator_models, seed_model_revision_file, seed_resource_names: ResourceNames
+    cognite_client: CogniteClient,
+    seed_simulator_models: SimulatorModel,
+    seed_model_revision_file: FileMetadata,
+    seed_resource_names: ResourceNames,
 ) -> None:
     model_unique_external_id = seed_resource_names.simulator_model_external_id
     model_revision_unique_external_id = seed_resource_names.simulator_model_revision_external_id
@@ -179,22 +188,24 @@ def seed_simulator_model_revisions(
 
 
 @pytest.fixture(scope="session")
-def seed_simulator_routines(cognite_client: CogniteClient, seed_simulator_model_revisions):
+def seed_simulator_routines(
+    cognite_client: CogniteClient, seed_simulator_model_revisions: None
+) -> Iterator[SimulatorRoutineList]:
     model_unique_external_id = RESOURCES.simulator_model_external_id
     simulator_routine_unique_external_id = RESOURCES.simulator_routine_external_id
 
     routines = cognite_client.simulators.routines.create(
         [
             SimulatorRoutineWrite.load(
                 {
-                    **SIMULATOR_ROUTINE,
+                    **SIMULATOR_ROUTINE.dump(),
                     "modelExternalId": model_unique_external_id,
                     "externalId": simulator_routine_unique_external_id,
                 }
             ),
             SimulatorRoutineWrite.load(
                 {
-                    **SIMULATOR_ROUTINE,
+                    **SIMULATOR_ROUTINE.dump(),
                     "modelExternalId": model_unique_external_id,
                     "externalId": simulator_routine_unique_external_id + "_1",
                 }
@@ -211,21 +222,24 @@ def seed_simulator_routines(cognite_client: CogniteClient, seed_simulator_model_
 
 def seed_simulator_routine_revision(
     cognite_client: CogniteClient, routine_external_id: str, version: str
-) -> dict[str, Any]:
+) -> SimulatorRoutineRevision:
     routine_revs = cognite_client.simulators.routines.revisions.list(routine_external_ids=[routine_external_id])
     rev_external_id = f"{routine_external_id}_{version}"
-    routine_rev_exists = routine_revs.get(external_id=rev_external_id)
+    revision = routine_revs.get(external_id=rev_external_id)
 
-    revision = {**SIMULATOR_ROUTINE_REVISION_OBJ, "externalId": rev_external_id}
-
-    if not routine_rev_exists:
-        cognite_client.simulators.routines.revisions.create(SimulatorRoutineRevisionWrite.load(revision))
+    if not revision:
+        revision_write = SimulatorRoutineRevisionWrite.load(
+            {**SIMULATOR_ROUTINE_REVISION.dump(), "externalId": rev_external_id}
+        )
+        return cognite_client.simulators.routines.revisions.create(revision_write)
 
     return revision
 
 
 @pytest.fixture(scope="session")
-def seed_simulator_routine_revisions(cognite_client: CogniteClient, seed_simulator_routines):
+def seed_simulator_routine_revisions(
+    cognite_client: CogniteClient, seed_simulator_routines: list[SimulatorRoutine]
+) -> Iterator[tuple[SimulatorRoutineRevision, SimulatorRoutineRevision]]:
     simulator_routine_external_id = RESOURCES.simulator_routine_external_id
 
     rev1 = seed_simulator_routine_revision(cognite_client, simulator_routine_external_id, "v1")