Skip to content

Commit

Permalink
add test for to_pandas for typed instances
Browse files Browse the repository at this point in the history
  • Loading branch information
haakonvt committed Sep 5, 2024
1 parent 89a9e01 commit 4cad8c5
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 2 deletions.
2 changes: 1 addition & 1 deletion cognite/client/data_classes/data_modeling/instances.py
Original file line number Diff line number Diff line change
Expand Up @@ -468,7 +468,7 @@ def to_pandas( # type: ignore [override]
from cognite.client.data_classes.data_modeling.typed_instances import _TypedInstance

if isinstance(self, _TypedInstance):
view_id, *extra = [self.get_source()] # type: ignore [attr-defined]
view_id, *extra = [self.get_source()]
else:
view_id, *extra = self.properties.keys()
# We only do/allow this if we have a single source:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,9 @@

from cognite.client.data_classes.data_modeling import DirectRelationReference, NodeList, ViewId
from cognite.client.data_classes.data_modeling.cdm.v1 import (
CogniteAsset, CogniteAssetApply, CogniteDescribableEdgeApply
CogniteAsset,
CogniteAssetApply,
CogniteDescribableEdgeApply,
)
from cognite.client.data_classes.data_modeling.typed_instances import (
PropertyOptions,
Expand Down Expand Up @@ -345,6 +347,8 @@ def test_cognite_asset_read_and_write(
xid = "externalId" if camel_case else "external_id"
assert xid in dumped
# Check that type/type_ are correctly dumped:
# - not starting with double underscore
# - not ending with single underscore
assert dumped["type"] == {"space": "should-be", xid: "at-root"}
if is_write:
properties = dumped["sources"][0]["properties"]
Expand All @@ -354,6 +358,41 @@ def test_cognite_asset_read_and_write(
# Check that properties are cased according to use_attribute_name:
assert ("source_created_time" if use_attribute_name else "sourceCreatedTime") in properties

@pytest.mark.dsl
@pytest.mark.parametrize("camel_case", (True, False))
def test_cognite_asset_read_and_write__to_pandas(
self, cognite_asset_kwargs: dict[str, Any], camel_case: bool
) -> None:
import pandas as pd

# When calling to_pandas, `use_attribute_name = not camel_case` due to how we expect
# attributes to be snake cased in python (in general).
asset_write = CogniteAssetApply(**cognite_asset_kwargs)
asset_read = CogniteAsset(**cognite_asset_kwargs, version=1, last_updated_time=10, created_time=5)

# TODO: Implement expand_properties=True for Apply classes?
write_df = asset_write.to_pandas(camel_case=camel_case)
read_df = asset_read.to_pandas(expand_properties=False, camel_case=camel_case)
read_expanded_df = asset_read.to_pandas(expand_properties=True, camel_case=camel_case)

xid = "externalId" if camel_case else "external_id"
for df in [write_df, read_df]:
assert df.index.is_unique
assert df.index[1] == xid
assert df.at["type", "value"] == {"space": "should-be", xid: "at-root"}

assert not read_expanded_df.index.is_unique # because 'type' is repeated
expected_type_df = pd.DataFrame(
[
({"space": "should-be", xid: "at-root"},),
({"space": "should-be", xid: "in-properties"},),
],
columns=["value"],
index=["type", "type"],
)
pd.testing.assert_frame_equal(read_expanded_df.loc["type"], expected_type_df)
assert ("sourceCreatedTime" if camel_case else "source_created_time") in read_expanded_df.index


@pytest.mark.parametrize(
"name, instance",
Expand Down

0 comments on commit 4cad8c5

Please sign in to comment.