Fix tests

datasets/flwr_datasets/partitioner/cid_partitioner.py

@@ -36,7 +36,7 @@ def _create_int_idx_to_cid(self) -> None:
 
         Client ids come from the columns specified in `partition_by`.
         """
-        unique_cid = self._dataset.unique(self._partition_by)
+        unique_cid = self.dataset.unique(self._partition_by)
         self._index_to_cid = dict(zip(range(len(unique_cid)), unique_cid))
 
     def _save_partition_indexing(self, idx: int, rows: List[int]) -> None:
@@ -64,7 +64,7 @@ def load_partition(self, idx: int) -> datasets.Dataset:
         if len(self._index_to_cid) == 0:
             self._create_int_idx_to_cid()
 
-        return self._dataset.filter(
+        return self.dataset.filter(
             lambda row: row[self._partition_by] == self._index_to_cid[idx]
         )
 
@@ -73,6 +73,7 @@ def index_to_cid(self) -> Dict[int, str]:
         """Index to corresponding cid from the dataset property."""
         return self._index_to_cid
 
+    # pylint: disable=R0201
     @index_to_cid.setter
     def index_to_cid(self, value: Dict[int, str]) -> None:
         raise AttributeError("Setting the index_to_cid dictionary is not allowed.")

datasets/flwr_datasets/partitioner/cid_partitioner_test.py

@@ -100,13 +100,13 @@ def test_correct_number_of_partitions(
         self, num_rows: int, num_unique_cid: int
     ) -> None:
         """Test if the # of available partitions is equal to # of unique clients."""
-        dataset, partitioner = _dummy_setup(num_rows, num_unique_cid)
+        _, partitioner = _dummy_setup(num_rows, num_unique_cid)
         _ = partitioner.load_partition(idx=0)
         self.assertEqual(len(partitioner.index_to_cid), num_unique_cid)
 
     def test_cannot_set_index_to_cid(self) -> None:
         """Test the lack of ability to set index_to_cid."""
-        dataset, partitioner = _dummy_setup(num_rows=10, n_unique_cids=2)
+        _, partitioner = _dummy_setup(num_rows=10, n_unique_cids=2)
         with self.assertRaises(AttributeError):
             partitioner.index_to_cid = {0: "0"}