From 39e2c9a83dc191d5cd4cc9a2cbb9d060f974f3f8 Mon Sep 17 00:00:00 2001 From: Karl Higley Date: Fri, 6 Jan 2023 12:22:02 -0500 Subject: [PATCH] Improve robustness of test for converting schemas to dataframes This test was brittle to the addition of new fields in the `ColumnSchema` dataclass, but minor rework avoids that issue. --- tests/unit/schema/test_schema.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/unit/schema/test_schema.py b/tests/unit/schema/test_schema.py index 1a648c257..239f89441 100644 --- a/tests/unit/schema/test_schema.py +++ b/tests/unit/schema/test_schema.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import dataclasses + import pytest from merlin.dag import ColumnSelector @@ -157,8 +159,11 @@ def test_schema_to_pandas(): schema_set = Schema(["a", "b", "c"]) df = schema_set.to_pandas() + expected_columns = [field.name for field in dataclasses.fields(ColumnSchema)] + expected_columns.remove("properties") + assert isinstance(df, pd.DataFrame) - assert list(df.columns) == ["name", "tags", "dtype", "is_list", "is_ragged"] + assert list(df.columns) == expected_columns def test_construct_schema_with_column_names():