improve tests more

cognite/client/utils/_auxiliary.py

@@ -16,6 +16,7 @@
 from urllib.parse import quote
 
 from cognite.client.utils import _json
+from cognite.client.utils._importing import import_yaml_with_converters
 from cognite.client.utils._text import (
     convert_all_keys_to_camel_case,
     convert_all_keys_to_snake_case,
@@ -83,7 +84,7 @@ def fast_dict_load(
 
 def load_yaml_or_json(resource: str) -> Any:
     try:
-        import yaml
+        yaml = import_yaml_with_converters()
 
         return yaml.safe_load(resource)
     except ImportError:

tests/tests_unit/test_base.py

@@ -8,7 +8,6 @@
 from unittest.mock import MagicMock
 
 import pytest
-import yaml
 
 from cognite.client import ClientConfig, CogniteClient
 from cognite.client.credentials import Token
@@ -306,8 +305,7 @@ def test_yaml_serialize(self, cognite_object_subclass: type[CogniteObject], cogn
             seed=65, cognite_client=cognite_mock_client_placeholder
         ).create_instance(cognite_object_subclass)
 
-        dumped = instance.dump(camel_case=True)
-        yaml_serialised = yaml.safe_dump(dumped)
+        yaml_serialised = instance.dump_yaml()
         loaded = instance.load(yaml_serialised, cognite_client=cognite_mock_client_placeholder)
 
         assert loaded.dump() == instance.dump()

tests/tests_unit/test_data_classes/test_data_models/test_queries.py

@@ -1,3 +1,4 @@
+import textwrap
 from collections.abc import Iterator
 from typing import Any
 
@@ -7,6 +8,7 @@
 from cognite.client.data_classes import filters as f
 from cognite.client.data_classes.data_modeling import InstanceSort, ViewId
 from cognite.client.data_classes.data_modeling import query as q
+from cognite.client.data_classes.data_modeling.cdm.v1 import CogniteFile
 
 
 def result_set_expression_load_and_dump_equals_data() -> Iterator[ParameterSet]:
@@ -144,37 +146,38 @@ def test_dump(self, raw_data: dict, loaded: q.Select) -> None:
 
 
 def query_load_yaml_data() -> Iterator[ParameterSet]:
-    raw_yaml = """with:
-    airplanes:
-        nodes:
-            filter:
-                equals:
-                    property: ["node", "externalId"]
-                    value: {"parameter": "airplaneExternalId"}
-            chainTo: destination
-            direction: outwards
-        limit: 1
-    lands_in_airports:
-        edges:
-            from: airplanes
-            maxDistance: 1
-            direction: outwards
-            filter:
-                equals:
-                    property: ["edge", "type"]
-                    value: ["aviation", "lands-in"]
-            chainTo: destination
-    airports:
-        nodes:
-            chainTo: destination
-            direction: outwards
-            from: lands_in_airports
-parameters:
-    airplaneExternalId: myFavouriteAirplane
-select:
-    airplanes: {}
-    airports: {}
-"""
+    raw_yaml = """\
+        with:
+            airplanes:
+                nodes:
+                    filter:
+                        equals:
+                            property: ["node", "externalId"]
+                            value: {"parameter": "airplaneExternalId"}
+                    chainTo: destination
+                    direction: outwards
+                limit: 1
+            lands_in_airports:
+                edges:
+                    from: airplanes
+                    maxDistance: 1
+                    direction: outwards
+                    filter:
+                        equals:
+                            property: ["edge", "type"]
+                            value: ["aviation", "lands-in"]
+                    chainTo: destination
+            airports:
+                nodes:
+                    chainTo: destination
+                    direction: outwards
+                    from: lands_in_airports
+        parameters:
+            airplaneExternalId: myFavouriteAirplane
+        select:
+            airplanes: {}
+            airports: {}
+    """
     expected = q.Query(
         with_={
             "airplanes": q.NodeResultSetExpression(
@@ -191,34 +194,35 @@ def query_load_yaml_data() -> Iterator[ParameterSet]:
         parameters={"airplaneExternalId": "myFavouriteAirplane"},
         select={"airplanes": q.Select(), "airports": q.Select()},
     )
-    yield pytest.param(raw_yaml, expected, id="Documentation Example")
-
-    raw_yaml = """with:
-  movies:
-    nodes:
-      filter:
-        equals:
-          property:
-          - IntegrationTestsImmutable
-          - Movie/2
-          - releaseYear
-          value: 1994
-      chainTo: destination
-      direction: outwards
-select:
-  movies:
-    sources:
-    - source:
-        space: IntegrationTestsImmutable
-        externalId: Movie
-        version: '2'
-        type: view
-      properties:
-      - title
-      - releaseYear
-cursors:
-  movies: Z0FBQUFBQmtwc0RxQmducHpsWFd6VnZFdWwyWnFJbmxWS1BlT
-"""
+    yield pytest.param(textwrap.dedent(raw_yaml), expected, id="Documentation Example")
+
+    raw_yaml = """\
+      with:
+        movies:
+          nodes:
+            filter:
+              equals:
+                property:
+                - IntegrationTestsImmutable
+                - Movie/2
+                - releaseYear
+                value: 1994
+            chainTo: destination
+            direction: outwards
+      select:
+        movies:
+          sources:
+          - source:
+              space: IntegrationTestsImmutable
+              externalId: Movie
+              version: '2'
+              type: view
+            properties:
+            - title
+            - releaseYear
+      cursors:
+        movies: Z0FBQUFBQmtwc0RxQmducHpsWFd6VnZFdWwyWnFJbmxWS1BlT
+    """
     movie_id = ViewId(space="IntegrationTestsImmutable", external_id="Movie", version="2")
     movies_released_1994 = q.NodeResultSetExpression(
         filter=f.Equals(list(movie_id.as_property_ref("releaseYear")), 1994)
@@ -228,11 +232,44 @@ def query_load_yaml_data() -> Iterator[ParameterSet]:
         select={"movies": q.Select([q.SourceSelector(movie_id, ["title", "releaseYear"])])},
         cursors={"movies": "Z0FBQUFBQmtwc0RxQmducHpsWFd6VnZFdWwyWnFJbmxWS1BlT"},
     )
-    yield pytest.param(raw_yaml, expected, id="Example with cursors")
+    yield pytest.param(textwrap.dedent(raw_yaml), expected, id="Example with cursors")
+
+
+@pytest.fixture
+def query_with_non_yaml_native_data_classes() -> q.Query:
+    # YAML doesn't support tuple, so we want to test that yaml tags e.g. !!python/tuple does not end
+    # up in the output when dumping a Query object.
+    return q.Query(
+        with_={
+            "files": q.NodeResultSetExpression(
+                filter=f.Exists(CogniteFile.get_source().as_property_ref("category")),
+                limit=None,
+            ),
+            "categories": q.NodeResultSetExpression(
+                from_="files",
+                through=CogniteFile.get_source().as_property_ref("category"),
+            ),
+        },
+        select={"categories": q.Select()},
+    )
 
 
 class TestQuery:
     @pytest.mark.parametrize("raw_data, expected", list(query_load_yaml_data()))
     def test_load_yaml(self, raw_data: str, expected: q.Query) -> None:
         actual = q.Query.load_yaml(raw_data)
         assert actual.dump(camel_case=True) == expected.dump(camel_case=True)
+
+    def test_dump_yaml_no_tags(self, query_with_non_yaml_native_data_classes: q.Query) -> None:
+        query = query_with_non_yaml_native_data_classes
+        dumped = query.dump_yaml()
+        assert "!!python/tuple" not in dumped
+
+        # Load will now load tuple as list:
+        loaded = q.Query.load_yaml(dumped)
+        assert loaded != query
+
+        # ...re-dump-loading should be equal to the first loaded
+        dumped_again = loaded.dump_yaml()
+        assert "!!python/tuple" not in dumped_again
+        assert q.Query.load_yaml(dumped_again) == loaded

tests/tests_unit/test_data_classes/test_hosted_extractors/test_jobs.py

@@ -1,23 +1,28 @@
+import textwrap
+
 import yaml
 
 from cognite.client.data_classes.hosted_extractors.jobs import Job
 
 
 class TestJob:
     def test_load_yaml_dump_unknown_config(self) -> None:
-        raw_yaml = """externalId: myJob
-sourceId: my_eventhub
-destinationId: EventHubTarget
-targetStatus: running
-status: running
-createdTime: 123
-lastUpdatedTime: 1234
-format:
-  type: value
-  encoding: utf16
-  compression: gzip
-config:
-  some: new_config
-  that: has not been seen before"""
-
+        raw_yaml = textwrap.dedent(
+            """\
+            externalId: myJob
+            sourceId: my_eventhub
+            destinationId: EventHubTarget
+            targetStatus: running
+            status: running
+            createdTime: 123
+            lastUpdatedTime: 1234
+            format:
+              type: value
+              encoding: utf16
+              compression: gzip
+            config:
+              some: new_config
+              that: has not been seen before
+            """
+        )
         assert Job.load(raw_yaml).dump() == yaml.safe_load(raw_yaml)

tests/tests_unit/test_utils/test_importing.py

@@ -1,4 +1,5 @@
 import pytest
+import yaml as normal_yaml
 
 from cognite.client.exceptions import CogniteImportError
 from cognite.client.utils._importing import import_yaml_with_converters, local_import
@@ -34,8 +35,7 @@ def test_dsl_deps_not_installed(self):
                 local_import(dep)
 
     def test_import_yaml_with_converters(self):
-        yaml = local_import("yaml")
-        assert yaml.dump(("foo", "bar")) == "!!python/tuple\n- foo\n- bar\n"
+        assert normal_yaml.dump(("foo", "bar")) == "!!python/tuple\n- foo\n- bar\n"
 
         better_yaml = import_yaml_with_converters()
         assert better_yaml.dump(("foo", "bar")) == "- foo\n- bar\n"