Skip to content

Commit

Permalink
Improve robustness of test for converting schemas to dataframes
Browse files Browse the repository at this point in the history
This test was brittle to the addition of new fields in the `ColumnSchema` dataclass, but minor rework avoids that issue.
  • Loading branch information
karlhigley committed Jan 6, 2023
1 parent 7fa799d commit 39e2c9a
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion tests/unit/schema/test_schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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():
Expand Down

0 comments on commit 39e2c9a

Please sign in to comment.