From 5b2c1ca58f0e524798d4b05f5813e292d47690d5 Mon Sep 17 00:00:00 2001 From: Adam Narozniak Date: Wed, 17 Jul 2024 12:59:48 +0200 Subject: [PATCH 1/2] Update tests for multiple partitioners --- datasets/flwr_datasets/federated_dataset_test.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/datasets/flwr_datasets/federated_dataset_test.py b/datasets/flwr_datasets/federated_dataset_test.py index f65aa6346f3a..e821cc359f86 100644 --- a/datasets/flwr_datasets/federated_dataset_test.py +++ b/datasets/flwr_datasets/federated_dataset_test.py @@ -144,9 +144,12 @@ def test_multiple_partitioners(self) -> None: dataset_test_partition0 = dataset_fds.load_partition(0, self.test_split) dataset = datasets.load_dataset(self.dataset_name) + expected_len = len(dataset[self.test_split]) // num_test_partitions + mod = len(dataset[self.test_split]) % num_test_partitions + expected_len += 1 if 0 < mod else 0 self.assertEqual( len(dataset_test_partition0), - len(dataset[self.test_split]) // num_test_partitions, + expected_len ) def test_no_need_for_split_keyword_if_one_partitioner(self) -> None: From 1995edb7501cec07eec6f8396b258bcccc01b8b2 Mon Sep 17 00:00:00 2001 From: Adam Narozniak Date: Wed, 17 Jul 2024 13:06:31 +0200 Subject: [PATCH 2/2] Fix formatting --- datasets/flwr_datasets/federated_dataset_test.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/datasets/flwr_datasets/federated_dataset_test.py b/datasets/flwr_datasets/federated_dataset_test.py index e821cc359f86..7c29a80929e0 100644 --- a/datasets/flwr_datasets/federated_dataset_test.py +++ b/datasets/flwr_datasets/federated_dataset_test.py @@ -147,10 +147,7 @@ def test_multiple_partitioners(self) -> None: expected_len = len(dataset[self.test_split]) // num_test_partitions mod = len(dataset[self.test_split]) % num_test_partitions expected_len += 1 if 0 < mod else 0 - self.assertEqual( - len(dataset_test_partition0), - expected_len - ) + self.assertEqual(len(dataset_test_partition0), expected_len) def test_no_need_for_split_keyword_if_one_partitioner(self) -> None: """Test if partitions got with and without split args are the same."""