Update test_managed_table_dataset.py

Signed-off-by: Sajid Alam <sajid_alam@mckinsey.com>
kedro-org · Jun 5, 2024 · 5304aa2 · 5304aa2
1 parent 497ba98
commit 5304aa2
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/kedro-datasets/tests/databricks/test_managed_table_dataset.py b/kedro-datasets/tests/databricks/test_managed_table_dataset.py
@@ -314,7 +314,7 @@ def test_save_schema_spark(
         assert subset_expected_df.exceptAll(saved_table).count() == 0
 
     def test_save_schema_pandas(
-        self, subset_pandas_df: pd.DataFrame, subset_expected_df: DataFrame
+            self, subset_pandas_df: pd.DataFrame, subset_expected_df: DataFrame
     ):
         unity_ds = ManagedTableDataset(
             database="test",
@@ -345,8 +345,8 @@ def test_save_schema_pandas(
             table="test_save_pd_schema",
             dataframe_type="pandas"
         )
-        saved_table = saved_ds.load().sort_values("name").reset_index(drop=True)
-        expected_table = subset_expected_df.sort_values("name").reset_index(drop=True)
+        saved_table = saved_ds.load().toPandas().sort_values("name").reset_index(drop=True)
+        expected_table = subset_expected_df.toPandas().sort_values("name").reset_index(drop=True)
         assert saved_table.equals(expected_table)
 
     def test_save_overwrite(
@@ -486,7 +486,7 @@ def test_load_pandas(self, sample_pandas_df: pd.DataFrame):
         pandas_ds = ManagedTableDataset(
             database="test", table="test_load_pandas", dataframe_type="pandas"
         )
-        pandas_df = pandas_ds.load().sort_values("name", ignore_index=True)
+        pandas_df = pandas_ds.load().toPandas().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(expected_df)