diff --git a/kedro-datasets/tests/databricks/test_managed_table_dataset.py b/kedro-datasets/tests/databricks/test_managed_table_dataset.py index 03a85d27e..015a1f776 100644 --- a/kedro-datasets/tests/databricks/test_managed_table_dataset.py +++ b/kedro-datasets/tests/databricks/test_managed_table_dataset.py @@ -343,9 +343,11 @@ def test_save_schema_pandas( saved_ds = ManagedTableDataset( database="test", table="test_save_pd_schema", + dataframe_type="pandas" ) - saved_table = saved_ds.load() - assert subset_expected_df.exceptAll(saved_table).count() == 0 + saved_table = saved_ds.load().sort_values("name").reset_index(drop=True) + expected_table = subset_expected_df.sort_values("name").reset_index(drop=True) + assert saved_table.equals(expected_table) def test_save_overwrite( self, sample_spark_df: DataFrame, append_spark_df: DataFrame @@ -485,7 +487,6 @@ def test_load_pandas(self, sample_pandas_df: pd.DataFrame): database="test", table="test_load_pandas", dataframe_type="pandas" ) pandas_df = pandas_ds.load().sort_values("name", ignore_index=True) + expected_df = sample_pandas_df.sort_values("name", ignore_index=True) - assert isinstance(pandas_df, pd.DataFrame) and pandas_df.equals( - sample_pandas_df - ) + assert isinstance(pandas_df, pd.DataFrame) and pandas_df.equals(expected_df)