From 543da44acb329d6e610a234a6c39fd569453c1b7 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Mon, 24 Jan 2022 16:03:36 -0800 Subject: [PATCH] add assert Signed-off-by: Yee Hing Tong --- tests/flytekit/unit/core/test_structured_dataset.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/flytekit/unit/core/test_structured_dataset.py b/tests/flytekit/unit/core/test_structured_dataset.py index 49a8f50015..4e4309e292 100644 --- a/tests/flytekit/unit/core/test_structured_dataset.py +++ b/tests/flytekit/unit/core/test_structured_dataset.py @@ -314,6 +314,8 @@ def test_to_python_value_without_incoming_columns(): lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt) sd = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, StructuredDataset) assert sd.metadata.structured_dataset_type.columns == [] + sub_df = sd.open(pd.DataFrame).all() + assert sub_df.shape[1] == 2 # should also work if subset type is just an annotated pd.DataFrame lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt)