From 262c059f2ef12441084149db28671cd7e382c357 Mon Sep 17 00:00:00 2001 From: Elena Khaustova Date: Tue, 7 Jan 2025 12:44:36 +0000 Subject: [PATCH] Added test_callable_save Signed-off-by: Elena Khaustova --- .../partitions/test_partitioned_dataset.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/kedro-datasets/tests/partitions/test_partitioned_dataset.py b/kedro-datasets/tests/partitions/test_partitioned_dataset.py index 921ae9774..01abf7dc1 100644 --- a/kedro-datasets/tests/partitions/test_partitioned_dataset.py +++ b/kedro-datasets/tests/partitions/test_partitioned_dataset.py @@ -52,6 +52,10 @@ def filepath_csvs(tmp_path): ] +def original_data_callable(): + return pd.DataFrame({"foo": 42, "bar": ["a", "b", None]}) + + class FakeDataset: # pylint: disable=too-few-public-methods pass @@ -101,6 +105,22 @@ def test_save(self, dataset, local_csvs, suffix): reloaded_data = loaded_partitions[part_id]() assert_frame_equal(reloaded_data, original_data) + @pytest.mark.parametrize("dataset", ["kedro_datasets.pickle.PickleDataset"]) + @pytest.mark.parametrize("suffix", ["", ".csv"]) + def test_callable_save(self, dataset, local_csvs, suffix): + pds = PartitionedDataset( + path=str(local_csvs), dataset=dataset, filename_suffix=suffix + ) + + part_id = "new/data" + pds.save({part_id: original_data_callable}) + + assert (local_csvs / "new" / ("data" + suffix)).is_file() + loaded_partitions = pds.load() + assert part_id in loaded_partitions + reloaded_data = loaded_partitions[part_id]() + assert reloaded_data == original_data_callable + @pytest.mark.parametrize("dataset", LOCAL_DATASET_DEFINITION) @pytest.mark.parametrize("suffix", ["", ".csv"]) def test_lazy_save(self, dataset, local_csvs, suffix):