test(datasets): use a more version-agnostic assert

Signed-off-by: Deepyaman Datta <deepyaman.datta@utexas.edu>
mark-druffel · Dec 5, 2024 · 10af4db · 10af4db
1 parent 6e411e4
commit 10af4db
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 6 deletions.
diff --git a/kedro-datasets/kedro_datasets/polars/csv_dataset.py b/kedro-datasets/kedro_datasets/polars/csv_dataset.py
@@ -52,15 +52,16 @@ class CSVDataset(AbstractVersionedDataset[pl.DataFrame, pl.DataFrame]):
 
     .. code-block:: pycon
 
-        >>> from kedro_datasets.polars import CSVDataset
         >>> import polars as pl
+        >>> from kedro_datasets.polars import CSVDataset
+        >>> from polars.testing import assert_frame_equal
         >>>
         >>> data = pl.DataFrame({"col1": [1, 2], "col2": [4, 5], "col3": [5, 6]})
         >>>
         >>> dataset = CSVDataset(filepath=tmp_path / "test.csv")
         >>> dataset.save(data)
         >>> reloaded = dataset.load()
-        >>> assert data.frame_equal(reloaded)
+        >>> assert_frame_equal(data, reloaded)
 
     """
 

diff --git a/kedro-datasets/kedro_datasets/polars/eager_polars_dataset.py b/kedro-datasets/kedro_datasets/polars/eager_polars_dataset.py
@@ -43,15 +43,16 @@ class EagerPolarsDataset(AbstractVersionedDataset[pl.DataFrame, pl.DataFrame]):
 
     .. code-block:: pycon
 
-        >>> from kedro_datasets.polars import EagerPolarsDataset
         >>> import polars as pl
+        >>> from kedro_datasets.polars import EagerPolarsDataset
+        >>> from polars.testing import assert_frame_equal
         >>>
         >>> data = pl.DataFrame({"col1": [1, 2], "col2": [4, 5], "col3": [5, 6]})
         >>>
         >>> dataset = EagerPolarsDataset(filepath=tmp_path / "test.parquet", file_format="parquet")
         >>> dataset.save(data)
         >>> reloaded = dataset.load()
-        >>> assert data.frame_equal(reloaded)
+        >>> assert_frame_equal(data, reloaded)
 
     """
 

diff --git a/kedro-datasets/kedro_datasets/polars/lazy_polars_dataset.py b/kedro-datasets/kedro_datasets/polars/lazy_polars_dataset.py
@@ -61,15 +61,16 @@ class LazyPolarsDataset(
 
     .. code-block:: pycon
 
-        >>> from kedro_datasets.polars import LazyPolarsDataset
         >>> import polars as pl
+        >>> from kedro_datasets.polars import LazyPolarsDataset
+        >>> from polars.testing import assert_frame_equal
         >>>
         >>> data = pl.DataFrame({"col1": [1, 2], "col2": [4, 5], "col3": [5, 6]})
         >>>
         >>> dataset = LazyPolarsDataset(filepath=tmp_path / "test.csv", file_format="csv")
         >>> dataset.save(data)
         >>> reloaded = dataset.load()
-        >>> assert data.frame_equal(reloaded.collect())
+        >>> assert_frame_equal(data, reloaded)
 
     """