From a2f9a1c5e0dbd0e681ea810e7eac65b0967eba5c Mon Sep 17 00:00:00 2001 From: Adam Narozniak Date: Mon, 26 Feb 2024 09:59:19 +0100 Subject: [PATCH] Test for empty partitions --- .../partitioner/shard_partitioner_test.py | 21 ++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/datasets/flwr_datasets/partitioner/shard_partitioner_test.py b/datasets/flwr_datasets/partitioner/shard_partitioner_test.py index 8aede604e188..47968699bba7 100644 --- a/datasets/flwr_datasets/partitioner/shard_partitioner_test.py +++ b/datasets/flwr_datasets/partitioner/shard_partitioner_test.py @@ -31,7 +31,7 @@ def _dummy_setup( shard_size: Optional[int], keep_incomplete_shard: bool = False, ) -> Tuple[Dataset, ShardPartitioner]: - """Create a dummy dataset for testing..""" + """Create a dummy dataset for testing.""" data = { partition_by: [i % 3 for i in range(num_rows)], "features": list(range(num_rows)), @@ -368,6 +368,25 @@ def test_incorrect_specification(self) -> None: with self.assertRaises(ValueError): _ = partitioner.load_partition(0) + def test_too_big_shard_size(self) -> None: + """Test if it is impossible to create an empty partition.""" + partition_by = "label" + num_rows = 20 + num_partitions = 3 + num_shards_per_node = None + shard_size = 10 + keep_incomplete_shard = False + _, partitioner = _dummy_setup( + num_rows, + partition_by, + num_partitions, + num_shards_per_node, + shard_size, + keep_incomplete_shard, + ) + with self.assertRaises(ValueError): + _ = partitioner.load_partition(2).num_rows + if __name__ == "__main__": unittest.main()