Repository: NVIDIA/NVFlare

Issue:
Build tests for new aggregation classes
- The new `InTimeAccumulateWeightedAggregator` now uses `WeightedAggregationHelper` and `DXOAggregator` which need to be unit tested.
- `InTimeAccumulateWeightedAggregator` also introduced support for DXO DataKind "COLLECTION" which needs to be tested.
- Also add tests for `exclude_vars` argument in `InTimeAccumulateWeightedAggregator`.

Discussion:

Patch:
diff --git a/nvflare/app_common/aggregators/dxo_aggregator.py b/nvflare/app_common/aggregators/dxo_aggregator.py
index 44dc3ce7..786c5a70 100644
--- a/nvflare/app_common/aggregators/dxo_aggregator.py
+++ b/nvflare/app_common/aggregators/dxo_aggregator.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+from typing import Any, Dict, Optional
 
 from nvflare.apis.dxo import DXO, DataKind, MetaKey
 from nvflare.apis.fl_component import FLComponent
@@ -24,18 +25,18 @@ from nvflare.app_common.app_constant import AppConstants
 class DXOAggregator(FLComponent):
     def __init__(
         self,
-        exclude_vars=None,
-        aggregation_weights=None,
+        exclude_vars: Optional[str] = None,
+        aggregation_weights: Optional[Dict[str, Any]] = None,
         expected_data_kind: DataKind = DataKind.WEIGHT_DIFF,
         name_postfix: str = "",
     ):
-        """Perform accumulated weighted aggregation for one kind of corresponding DXO from contributors
+        """Perform accumulated weighted aggregation for one kind of corresponding DXO from contributors.
 
         Args:
-            exclude_vars ([type], optional): regex to match excluded vars during aggregation. Defaults to None.
-            aggregation_weights ([type], optional): dictionary to map contributor name to its aggregation weights.
+            exclude_vars (str, optional): Regex to match excluded vars during aggregation. Defaults to None.
+            aggregation_weights (Dict[str, Any], optional): Aggregation weight for each contributor.
                                 Defaults to None.
-            expected_data_kind: Expected DataKind for this kind of DXO.
+            expected_data_kind (DataKind): Expected DataKind for this DXO.
             name_postfix: optional postfix to give to class name and show in logger output.
         """
         super().__init__()
diff --git a/nvflare/app_common/aggregators/intime_accumulate_model_aggregator.py b/nvflare/app_common/aggregators/intime_accumulate_model_aggregator.py
index dc49cc6c..1ceee909 100644
--- a/nvflare/app_common/aggregators/intime_accumulate_model_aggregator.py
+++ b/nvflare/app_common/aggregators/intime_accumulate_model_aggregator.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Any, Dict, Union
+
 from nvflare.apis.dxo import DXO, DataKind, from_shareable
 from nvflare.apis.fl_constant import ReservedKey, ReturnCode
 from nvflare.apis.fl_context import FLContext
@@ -21,22 +23,48 @@ from nvflare.app_common.aggregators.dxo_aggregator import DXOAggregator
 from nvflare.app_common.app_constant import AppConstants
 
 
+def _is_nested_aggregation_weights(aggregation_weights):
+    if not aggregation_weights:
+        return False
+    if not isinstance(aggregation_weights, dict):
+        return False
+    first_value = next(iter(aggregation_weights.items()))[1]
+    if not isinstance(first_value, dict):
+        return False
+    return True
+
+
+def _get_missing_keys(ref_dict: dict, dict_to_check: dict):
+    result = []
+    for k in ref_dict:
+        if k not in dict_to_check:
+            result.append(k)
+    return result
+
+
 class InTimeAccumulateWeightedAggregator(Aggregator):
-    def __init__(self, exclude_vars=None, aggregation_weights=None, expected_data_kind=DataKind.WEIGHT_DIFF):
+    def __init__(
+        self,
+        exclude_vars: Union[str, Dict[str, str], None] = None,
+        aggregation_weights: Union[Dict[str, Any], Dict[str, Dict[str, Any]], None] = None,
+        expected_data_kind: Union[DataKind, Dict[str, DataKind]] = DataKind.WEIGHT_DIFF,
+    ):
         """Perform accumulated weighted aggregation
         It parses the shareable and aggregates the contained DXO(s).
 
         Args:
-            exclude_vars ([type], optional): regex to match excluded vars during aggregation. Defaults to None.
-                                Can be one string or a dict of keys with regex strings corresponding to each aggregated
-                                DXO when processing a DXO of `DataKind.COLLECTION`.
-            aggregation_weights ([type], optional): dictionary to map contributor name to its aggregation weights.
-                                Defaults to None.
-                                Can be one dict or a dict of dicts corresponding to each aggregated DXO
-                                when processing DXO of `DataKind.COLLECTION`.
-            expected_data_kind: DataKind or dict of keys and matching DataKind entries
-                                when processing DXO of `DataKind.COLLECTION`.
-                                Only the keys in the dict will be processed.
+            exclude_vars (Union[str, Dict[str, str]], optional):
+                Regular expression string to match excluded vars during aggregation. Defaults to None.
+                Can be one string or a dict of {dxo_name: regex strings} corresponding to each aggregated DXO
+                when processing a DXO of `DataKind.COLLECTION`.
+            aggregation_weights (Union[Dict[str, Any], Dict[str, Dict[str, Any]]], optional):
+                Aggregation weight for each contributor. Defaults to None.
+                Can be one dict of {contrib_name: aggr_weight} or a dict of dicts corresponding to each aggregated DXO
+                when processing a DXO of `DataKind.COLLECTION`.
+            expected_data_kind (Union[DataKind, Dict[str, DataKind]]):
+                DataKind for DXO. Defaults to DataKind.WEIGHT_DIFF
+                Can be one DataKind or a dict of {dxo_name: DataKind} corresponding to each aggregated DXO
+                when processing a DXO of `DataKind.COLLECTION`. Only the keys in this dict will be processed.
         """
         super().__init__()
         self.logger.debug(f"exclude vars: {exclude_vars}")
@@ -49,22 +77,38 @@ class InTimeAccumulateWeightedAggregator(Aggregator):
         if isinstance(expected_data_kind, dict):
             for k, v in expected_data_kind.items():
                 if v not in [DataKind.WEIGHT_DIFF, DataKind.WEIGHTS]:
-                    raise ValueError(f"expected_data_kind[{k}] = {v} not in WEIGHT_DIFF or WEIGHTS")
+                    raise ValueError(
+                        f"expected_data_kind[{k}] = {v} is not {DataKind.WEIGHT_DIFF} or {DataKind.WEIGHTS}"
+                    )
             self.expected_data_kind = expected_data_kind
         else:
             if expected_data_kind not in [DataKind.WEIGHT_DIFF, DataKind.WEIGHTS]:
-                raise ValueError(f"expected_data_kind={expected_data_kind} not in WEIGHT_DIFF or WEIGHTS")
+                raise ValueError(
+                    f"expected_data_kind = {expected_data_kind} is not {DataKind.WEIGHT_DIFF} or {DataKind.WEIGHTS}"
+                )
             self.expected_data_kind = {self._single_dxo_key: expected_data_kind}
 
         # Check exclude_vars
+        if exclude_vars:
+            if not isinstance(exclude_vars, dict) and not isinstance(exclude_vars, str):
+                raise ValueError(
+                    f"exclude_vars = {exclude_vars} should be a regex string but got {type(exclude_vars)}."
+                )
+            if isinstance(exclude_vars, dict):
+                missing_keys = _get_missing_keys(expected_data_kind, exclude_vars)
+                if len(missing_keys) != 0:
+                    raise ValueError(
+                        "A dict exclude_vars should specify exclude_vars for every key in expected_data_kind. "
+                        f"But missed these keys: {missing_keys}"
+                    )
+
         exclude_vars_dict = dict()
         for k in self.expected_data_kind.keys():
             if isinstance(exclude_vars, dict):
                 if k in exclude_vars:
                     if not isinstance(exclude_vars[k], str):
                         raise ValueError(
-                            f"exclude_vars[{k}] = {exclude_vars[k]} not a string but {type(exclude_vars[k])}! "
-                            f"Expected type regex string."
+                            f"exclude_vars[{k}] = {exclude_vars[k]} should be a regex string but got {type(exclude_vars[k])}."
                         )
                     exclude_vars_dict[k] = exclude_vars[k]
             else:
@@ -75,6 +119,14 @@ class InTimeAccumulateWeightedAggregator(Aggregator):
         self.exclude_vars = exclude_vars_dict
 
         # Check aggregation weights
+        if _is_nested_aggregation_weights(aggregation_weights):
+            missing_keys = _get_missing_keys(expected_data_kind, aggregation_weights)
+            if len(missing_keys) != 0:
+                raise ValueError(
+                    "A dict of dict aggregation_weights should specify aggregation_weights "
+                    f"for every key in expected_data_kind. But missed these keys: {missing_keys}"
+                )
+
         aggregation_weights = aggregation_weights or {}
         aggregation_weights_dict = dict()
         for k in self.expected_data_kind.keys():
@@ -83,8 +135,6 @@ class InTimeAccumulateWeightedAggregator(Aggregator):
             else:
                 # assume same aggregation weights for each entry of DXO collection.
                 aggregation_weights_dict[k] = aggregation_weights
-        if self._single_dxo_key in self.expected_data_kind:
-            aggregation_weights_dict[self._single_dxo_key] = aggregation_weights
         self.aggregation_weights = aggregation_weights_dict
 
         # Set up DXO aggregators
diff --git a/nvflare/app_common/aggregators/weighted_aggregation_helper.py b/nvflare/app_common/aggregators/weighted_aggregation_helper.py
index d2cc51bf..1376d876 100644
--- a/nvflare/app_common/aggregators/weighted_aggregation_helper.py
+++ b/nvflare/app_common/aggregators/weighted_aggregation_helper.py
@@ -13,20 +13,15 @@
 # limitations under the License.
 
 import re
+from typing import Optional
 
 
 class WeightedAggregationHelper(object):
-    def __init__(self, exclude_vars=None):
-        """Perform weighted aggregation
-        It computes
-        weighted_sum = sum(dict*weight)
-
-        in accept function
-        The aggregate function returns
-        weighted_sum / sum_of_weights
+    def __init__(self, exclude_vars: Optional[str] = None):
+        """Perform weighted aggregation.
 
         Args:
-            exclude_vars ([type], optional): regex to match excluded vars during aggregation. Defaults to None.
+            exclude_vars (str, optional): regex string to match excluded vars during aggregation. Defaults to None.
         """
         super().__init__()
         self.exclude_vars = re.compile(exclude_vars) if exclude_vars else None
@@ -40,7 +35,8 @@ class WeightedAggregationHelper(object):
         self.counts = {}
         self.history = []
 
-    def add(self, data, weight, contributor_name, contribution_round) -> bool:
+    def add(self, data, weight, contributor_name, contribution_round):
+        """Compute weighted sum and sum of weights."""
         for k, v in data.items():
             if self.exclude_vars is not None and self.exclude_vars.search(k):
                 continue
@@ -61,6 +57,7 @@ class WeightedAggregationHelper(object):
         )
 
     def get_result(self):
+        """Divide weighted sum by sum of weights."""
         aggregated_dict = {k: v / self.counts[k] for k, v in self.total.items()}
         self.reset_stats()
         return aggregated_dict


Test patch:
diff --git a/test/unit_test/app_common/test_accumulate_weighted_aggregator.py b/test/unit_test/app_common/test_accumulate_weighted_aggregator.py
index 8571fc10..5148ed98 100644
--- a/test/unit_test/app_common/test_accumulate_weighted_aggregator.py
+++ b/test/unit_test/app_common/test_accumulate_weighted_aggregator.py
@@ -25,39 +25,40 @@ from nvflare.app_common.aggregators.accumulate_model_aggregator import Accumulat
 from nvflare.app_common.app_constant import AppConstants
 
 
-class TestAggregator:
-    @pytest.mark.parametrize("aggregator", [AccumulateWeightedAggregator])
+class TestAccumulateWeightedAggregator:
     @pytest.mark.parametrize(
         "received,expected",
         [
-            [
+            (
                 {"client1": {"weight": 0.5, "iter_number": 1, "aggr_data": {"var1": np.array([2.0, 3.0, 1.1, 0.1])}}},
                 {"var1": np.array([1.0, 1.5, 0.55, 0.05])},
-            ],
-            [
+            ),
+            (
                 {"client1": {"weight": 1.0, "iter_number": 1, "aggr_data": {"var1": np.array([2.0, 3.0, 1.1, 0.1])}}},
                 {"var1": np.array([2.0, 3.0, 1.1, 0.1])},
-            ],
-            [
+            ),
+            (
                 {
                     "client1": {"weight": 0.5, "iter_number": 3, "aggr_data": {"var1": np.array([2.0, 3.0, 1.1, 0.1])}},
                     "client2": {"weight": 1.0, "iter_number": 1, "aggr_data": {"var1": np.array([1.0, 1.0, 2.1, 0.5])}},
                 },
                 {"var1": np.array([1.0, 1.375, 0.9375, 0.1625])},
-            ],
-            [
+            ),
+            (
                 {
                     "client1": {"weight": 1.0, "iter_number": 5, "aggr_data": {"var1": np.array([2.0, 3.0, 1.1, 0.1])}},
                     "client2": {"weight": 0.5, "iter_number": 4, "aggr_data": {"var1": np.array([1.0, 1.0, 2.1, 0.5])}},
                 },
                 {"var1": np.array([4.0 / 3, 17.0 / 9, 9.7 / 9, 1.0 / 6])},
-            ],
+            ),
         ],
     )
-    def test_accum_aggregator(self, aggregator, received, expected):
+    def test_aggregate(self, received, expected):
         aggregation_weights = {k: v["weight"] for k, v in received.items()}
-        agg = aggregator(aggregation_weights=aggregation_weights)
+        agg = AccumulateWeightedAggregator(aggregation_weights=aggregation_weights)
         fl_ctx_mgr = FLContextManager(engine=None, identity_name="", run_num=1, public_stickers={}, private_stickers={})
+        fl_ctx = FLContext()
+        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
         for k, v in received.items():
             fl_ctx_mgr.identity_name = k
             dxo = DXO(
@@ -68,8 +69,6 @@ class TestAggregator:
                 },
             )
 
-            fl_ctx = FLContext()
-            fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
             s = Shareable()
             s.set_peer_props({ReservedKey.IDENTITY_NAME: k})
             s.set_header(AppConstants.CONTRIBUTION_ROUND, 0)
@@ -78,21 +77,20 @@ class TestAggregator:
         result = agg.aggregate(fl_ctx)
         np.testing.assert_allclose(result["DXO"]["data"]["var1"], expected["var1"])
 
-    @pytest.mark.parametrize("aggregator", [AccumulateWeightedAggregator])
-    @pytest.mark.parametrize("shape", [(4), (6, 6)])
+    @pytest.mark.parametrize("shape", [4, (6, 6)])
     @pytest.mark.parametrize("n_clients", [10, 50, 100])
-    def test_accum_aggregator_random(self, aggregator, shape, n_clients):
+    def test_aggregate_random(self, shape, n_clients):
         aggregation_weights = {f"client_{i}": random.random() for i in range(n_clients)}
-        agg = aggregator(aggregation_weights=aggregation_weights)
+        agg = AccumulateWeightedAggregator(aggregation_weights=aggregation_weights)
         weighted_sum = np.zeros(shape)
         sum_of_weights = 0
         fl_ctx_mgr = FLContextManager(engine=None, identity_name="", run_num=1, public_stickers={}, private_stickers={})
+        fl_ctx = FLContext()
+        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
         for client_name in aggregation_weights:
             iter_number = random.randint(1, 50)
             fl_ctx_mgr.identity_name = client_name
             weights = np.random.random(shape)
-            fl_ctx = FLContext()
-            fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
             s = Shareable()
             s.set_peer_props({ReservedKey.IDENTITY_NAME: client_name})
             s.set_header(AppConstants.CONTRIBUTION_ROUND, 0)
@@ -111,34 +109,21 @@ class TestAggregator:
         result_dxo = from_shareable(result)
         np.testing.assert_allclose(result_dxo.data["var1"], weighted_sum / sum_of_weights)
 
-    @pytest.mark.parametrize("aggregator", [AccumulateWeightedAggregator])
-    def test_accum_aggregator_accept(self, aggregator):
+    @pytest.mark.parametrize("current_round,contribution_round,expected", [(1, 1, True), (2, 1, False)])
+    def test_accept(self, current_round, contribution_round, expected):
         aggregation_weights = {f"client_{i}": random.random() for i in range(2)}
-        agg = aggregator(aggregation_weights=aggregation_weights)
+        agg = AccumulateWeightedAggregator(aggregation_weights=aggregation_weights)
         client_name = "client_0"
         iter_number = 1
         fl_ctx_mgr = FLContextManager(engine=None, identity_name="", run_num=1, public_stickers={}, private_stickers={})
         fl_ctx_mgr.identity_name = client_name
-        weights = np.random.random((4))
+        weights = np.random.random(4)
 
         fl_ctx = FLContext()
         s = Shareable()
         s.set_peer_props({ReservedKey.IDENTITY_NAME: client_name})
-        s.set_header(AppConstants.CONTRIBUTION_ROUND, 1)
-        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 1)
-        dxo = DXO(
-            DataKind.WEIGHT_DIFF,
-            data={"var1": weights},
-            meta={
-                MetaKey.NUM_STEPS_CURRENT_ROUND: iter_number,
-            },
-        )
-        assert (True) == agg.accept(dxo.update_shareable(s), fl_ctx)
-
-        client_name = "client_0"
-        iter_number = 1
-        fl_ctx_mgr.identity_name = client_name
-        weights = np.random.random((4))
+        s.set_header(AppConstants.CONTRIBUTION_ROUND, contribution_round)
+        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, current_round)
         dxo = DXO(
             DataKind.WEIGHT_DIFF,
             data={"var1": weights},
@@ -146,9 +131,4 @@ class TestAggregator:
                 MetaKey.NUM_STEPS_CURRENT_ROUND: iter_number,
             },
         )
-        fl_ctx = FLContext()
-        s = Shareable()
-        s.set_peer_props({ReservedKey.IDENTITY_NAME: client_name})
-        s.set_header(AppConstants.CONTRIBUTION_ROUND, 1)
-        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 2)
-        assert (False) == agg.accept(dxo.update_shareable(s), fl_ctx)
+        assert agg.accept(dxo.update_shareable(s), fl_ctx) == expected
diff --git a/test/unit_test/app_common/test_in_time_accumulate_weighted_aggregator.py b/test/unit_test/app_common/test_in_time_accumulate_weighted_aggregator.py
index 6ab83515..eb3413e9 100644
--- a/test/unit_test/app_common/test_in_time_accumulate_weighted_aggregator.py
+++ b/test/unit_test/app_common/test_in_time_accumulate_weighted_aggregator.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import random
+import re
 
 import numpy as np
 import pytest
@@ -25,20 +26,138 @@ from nvflare.app_common.aggregators.intime_accumulate_model_aggregator import In
 from nvflare.app_common.app_constant import AppConstants
 
 
-class TestAggregator:
-    @pytest.mark.parametrize("aggregator", [InTimeAccumulateWeightedAggregator])
+class TestInTimeAccumulateWeightedAggregator:
+    @pytest.mark.parametrize(
+        "exclude_vars,aggregation_weights,expected_data_kind,error,error_msg",
+        [
+            (
+                2.0,
+                None,
+                DataKind.WEIGHT_DIFF,
+                ValueError,
+                f"exclude_vars = 2.0 should be a regex string but got {type(2.0)}.",
+            ),
+            (
+                {"dxo1": 3.0, "dxo2": ""},
+                None,
+                {"dxo1": DataKind.WEIGHT_DIFF, "dxo2": DataKind.WEIGHT_DIFF},
+                ValueError,
+                f"exclude_vars[dxo1] = 3.0 should be a regex string but got {type(3.0)}.",
+            ),
+            (None, None, DataKind.ANALYTIC, ValueError, "expected_data_kind = ANALYTIC is not WEIGHT_DIFF or WEIGHTS"),
+            (
+                None,
+                None,
+                {"dxo1": DataKind.WEIGHT_DIFF, "dxo2": DataKind.ANALYTIC},
+                ValueError,
+                "expected_data_kind[dxo2] = ANALYTIC is not WEIGHT_DIFF or WEIGHTS",
+            ),
+            (
+                None,
+                {"dxo1": {"client_0": 1.0, "client_1": 2.0}},
+                {"dxo1": DataKind.WEIGHT_DIFF, "dxo2": DataKind.WEIGHT_DIFF},
+                ValueError,
+                "A dict of dict aggregation_weights should specify aggregation_weights "
+                "for every key in expected_data_kind. But missed these keys: ['dxo2']",
+            ),
+            (
+                {"dxo2": ""},
+                None,
+                {"dxo1": DataKind.WEIGHT_DIFF, "dxo2": DataKind.WEIGHT_DIFF},
+                ValueError,
+                "A dict exclude_vars should specify exclude_vars for every key in expected_data_kind. "
+                "But missed these keys: ['dxo1']",
+            ),
+        ],
+    )
+    def test_invalid_create(self, exclude_vars, aggregation_weights, expected_data_kind, error, error_msg):
+        with pytest.raises(error, match=re.escape(error_msg)):
+            _ = InTimeAccumulateWeightedAggregator(
+                exclude_vars=exclude_vars,
+                aggregation_weights=aggregation_weights,
+                expected_data_kind=expected_data_kind,
+            )
+
+    @pytest.mark.parametrize(
+        "exclude_vars,aggregation_weights,expected_data_kind,expected_object",
+        [
+            (
+                None,
+                None,
+                DataKind.WEIGHTS,
+                InTimeAccumulateWeightedAggregator(
+                    exclude_vars=None, aggregation_weights=None, expected_data_kind=DataKind.WEIGHTS
+                ),
+            ),
+            (
+                "hello",
+                None,
+                {"dxo1": DataKind.WEIGHTS, "dxo2": DataKind.WEIGHT_DIFF},
+                InTimeAccumulateWeightedAggregator(
+                    exclude_vars={"dxo1": "hello", "dxo2": "hello"},
+                    aggregation_weights=None,
+                    expected_data_kind={"dxo1": DataKind.WEIGHTS, "dxo2": DataKind.WEIGHT_DIFF},
+                ),
+            ),
+            (
+                None,
+                {"client_0": 1.0, "client_1": 2.0},
+                {"dxo1": DataKind.WEIGHTS, "dxo2": DataKind.WEIGHT_DIFF},
+                InTimeAccumulateWeightedAggregator(
+                    exclude_vars=None,
+                    aggregation_weights={
+                        "dxo1": {"client_0": 1.0, "client_1": 2.0},
+                        "dxo2": {"client_0": 1.0, "client_1": 2.0},
+                    },
+                    expected_data_kind={"dxo1": DataKind.WEIGHTS, "dxo2": DataKind.WEIGHT_DIFF},
+                ),
+            ),
+        ],
+    )
+    def test_create(self, exclude_vars, aggregation_weights, expected_data_kind, expected_object):
+        result = InTimeAccumulateWeightedAggregator(
+            exclude_vars=exclude_vars, aggregation_weights=aggregation_weights, expected_data_kind=expected_data_kind
+        )
+        assert result.exclude_vars == expected_object.exclude_vars
+        assert result.aggregation_weights == expected_object.aggregation_weights
+        assert result.expected_data_kind == expected_object.expected_data_kind
+
+    @pytest.mark.parametrize("current_round,contribution_round,expected", [(1, 1, True), (2, 1, False)])
+    def test_accept(self, current_round, contribution_round, expected):
+        aggregation_weights = {f"client_{i}": random.random() for i in range(2)}
+        agg = InTimeAccumulateWeightedAggregator(aggregation_weights=aggregation_weights)
+        client_name = "client_0"
+        iter_number = 1
+        fl_ctx_mgr = FLContextManager(engine=None, identity_name="", run_num=1, public_stickers={}, private_stickers={})
+        fl_ctx_mgr.identity_name = client_name
+        weights = np.random.random(4)
+
+        fl_ctx = FLContext()
+        s = Shareable()
+        s.set_peer_props({ReservedKey.IDENTITY_NAME: client_name})
+        s.set_header(AppConstants.CONTRIBUTION_ROUND, contribution_round)
+        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, current_round)
+        dxo = DXO(
+            DataKind.WEIGHT_DIFF,
+            data={"var1": weights},
+            meta={
+                MetaKey.NUM_STEPS_CURRENT_ROUND: iter_number,
+            },
+        )
+        assert agg.accept(dxo.update_shareable(s), fl_ctx) == expected
+
     @pytest.mark.parametrize(
         "received,expected",
         [
-            [
+            (
                 {"client1": {"weight": 0.5, "iter_number": 1, "aggr_data": {"var1": np.array([2.0, 3.0, 1.1, 0.1])}}},
                 {"var1": np.array([2.0, 3.0, 1.1, 0.1])},
-            ],
-            [
+            ),
+            (
                 {"client1": {"weight": 1.0, "iter_number": 1, "aggr_data": {"var1": np.array([2.0, 3.0, 1.1, 0.1])}}},
                 {"var1": np.array([2.0, 3.0, 1.1, 0.1])},
-            ],
-            [
+            ),
+            (
                 {
                     "client1": {"weight": 0.5, "iter_number": 1, "aggr_data": {"var1": np.array([2.0, 3.0, 1.1, 0.1])}},
                     "client2": {"weight": 1.0, "iter_number": 1, "aggr_data": {"var1": np.array([1.0, 1.0, 2.1, 0.5])}},
@@ -53,8 +172,8 @@ class TestAggregator:
                         ]
                     )
                 },
-            ],
-            [
+            ),
+            (
                 {
                     "client1": {"weight": 1.0, "iter_number": 2, "aggr_data": {"var1": np.array([2.0, 3.0, 1.1, 0.1])}},
                     "client2": {"weight": 1.0, "iter_number": 4, "aggr_data": {"var1": np.array([1.0, 1.0, 2.1, 0.5])}},
@@ -69,13 +188,15 @@ class TestAggregator:
                         ]
                     )
                 },
-            ],
+            ),
         ],
     )
-    def test_accum_aggregator(self, aggregator, received, expected):
+    def test_aggregate(self, received, expected):
         aggregation_weights = {k: v["weight"] for k, v in received.items()}
-        agg = aggregator(aggregation_weights=aggregation_weights)
+        agg = InTimeAccumulateWeightedAggregator(aggregation_weights=aggregation_weights)
         fl_ctx_mgr = FLContextManager(engine=None, identity_name="", run_num=1, public_stickers={}, private_stickers={})
+        fl_ctx = FLContext()
+        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
         for k, v in received.items():
             fl_ctx_mgr.identity_name = k
             dxo = DXO(
@@ -86,8 +207,6 @@ class TestAggregator:
                 },
             )
 
-            fl_ctx = FLContext()
-            fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
             s = Shareable()
             s.set_peer_props({ReservedKey.IDENTITY_NAME: k})
             s.set_header(AppConstants.CONTRIBUTION_ROUND, 0)
@@ -96,21 +215,20 @@ class TestAggregator:
         result = agg.aggregate(fl_ctx)
         np.testing.assert_allclose(result["DXO"]["data"]["var1"], expected["var1"])
 
-    @pytest.mark.parametrize("aggregator", [InTimeAccumulateWeightedAggregator])
-    @pytest.mark.parametrize("shape", [(4), (6, 6)])
+    @pytest.mark.parametrize("shape", [4, (6, 6)])
     @pytest.mark.parametrize("n_clients", [10, 50, 100])
-    def test_accum_aggregator_random(self, aggregator, shape, n_clients):
+    def test_aggregate_random(self, shape, n_clients):
         aggregation_weights = {f"client_{i}": random.random() for i in range(n_clients)}
-        agg = aggregator(aggregation_weights=aggregation_weights)
+        agg = InTimeAccumulateWeightedAggregator(aggregation_weights=aggregation_weights)
         weighted_sum = np.zeros(shape)
         sum_of_weights = 0
         fl_ctx_mgr = FLContextManager(engine=None, identity_name="", run_num=1, public_stickers={}, private_stickers={})
+        fl_ctx = FLContext()
+        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
         for client_name in aggregation_weights:
             iter_number = random.randint(1, 50)
             fl_ctx_mgr.identity_name = client_name
             weights = np.random.random(shape)
-            fl_ctx = FLContext()
-            fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
             s = Shareable()
             s.set_peer_props({ReservedKey.IDENTITY_NAME: client_name})
             s.set_header(AppConstants.CONTRIBUTION_ROUND, 0)
@@ -129,44 +247,55 @@ class TestAggregator:
         result_dxo = from_shareable(result)
         np.testing.assert_allclose(result_dxo.data["var1"], weighted_sum / sum_of_weights)
 
-    @pytest.mark.parametrize("aggregator", [InTimeAccumulateWeightedAggregator])
-    def test_accum_aggregator_accept(self, aggregator):
-        aggregation_weights = {f"client_{i}": random.random() for i in range(2)}
-        agg = aggregator(aggregation_weights=aggregation_weights)
-        client_name = "client_0"
-        iter_number = 1
+    @pytest.mark.parametrize("num_dxo", [1, 2, 3])
+    @pytest.mark.parametrize("shape", [4, (6, 6)])
+    @pytest.mark.parametrize("n_clients", [10, 50, 100])
+    def test_aggregate_random_dxos(self, num_dxo, shape, n_clients):
+        dxo_names = [f"dxo_{i}" for i in range(num_dxo)]
+        client_names = [f"client_{i}" for i in range(n_clients)]
+        aggregation_weights = {
+            dxo_name: {client_name: random.random() for client_name in client_names} for dxo_name in dxo_names
+        }
+        agg = InTimeAccumulateWeightedAggregator(
+            aggregation_weights=aggregation_weights,
+            expected_data_kind={dxo_name: DataKind.WEIGHT_DIFF for dxo_name in dxo_names},
+        )
+        weighted_sum = {dxo_name: np.zeros(shape) for dxo_name in dxo_names}
+        sum_of_weights = {dxo_name: 0 for dxo_name in dxo_names}
         fl_ctx_mgr = FLContextManager(engine=None, identity_name="", run_num=1, public_stickers={}, private_stickers={})
-        fl_ctx_mgr.identity_name = client_name
-        weights = np.random.random((4))
-
         fl_ctx = FLContext()
-        s = Shareable()
-        s.set_peer_props({ReservedKey.IDENTITY_NAME: client_name})
-        s.set_header(AppConstants.CONTRIBUTION_ROUND, 1)
-        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 1)
-        dxo = DXO(
-            DataKind.WEIGHT_DIFF,
-            data={"var1": weights},
-            meta={
-                MetaKey.NUM_STEPS_CURRENT_ROUND: iter_number,
-            },
-        )
-        assert (True) == agg.accept(dxo.update_shareable(s), fl_ctx)
+        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
+        for client_name in client_names:
+            iter_number = random.randint(1, 50)
+            fl_ctx_mgr.identity_name = client_name
 
-        client_name = "client_0"
-        iter_number = 1
-        fl_ctx_mgr.identity_name = client_name
-        weights = np.random.random((4))
-        dxo = DXO(
-            DataKind.WEIGHT_DIFF,
-            data={"var1": weights},
-            meta={
-                MetaKey.NUM_STEPS_CURRENT_ROUND: iter_number,
-            },
-        )
-        fl_ctx = FLContext()
-        s = Shareable()
-        s.set_peer_props({ReservedKey.IDENTITY_NAME: client_name})
-        s.set_header(AppConstants.CONTRIBUTION_ROUND, 1)
-        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 2)
-        assert (False) == agg.accept(dxo.update_shareable(s), fl_ctx)
+            dxo_collection_data = {}
+            for dxo_name in dxo_names:
+                values = np.random.random(shape)
+                dxo = DXO(
+                    data_kind=DataKind.WEIGHT_DIFF,
+                    data={"var1": values},
+                    meta={
+                        MetaKey.NUM_STEPS_CURRENT_ROUND: iter_number,
+                    },
+                )
+                dxo_collection_data[dxo_name] = dxo
+                weighted_sum[dxo_name] = (
+                    weighted_sum[dxo_name] + values * iter_number * aggregation_weights[dxo_name][client_name]
+                )
+                sum_of_weights[dxo_name] = (
+                    sum_of_weights[dxo_name] + iter_number * aggregation_weights[dxo_name][client_name]
+                )
+
+            dxo_collection = DXO(data_kind=DataKind.COLLECTION, data=dxo_collection_data)
+            s = Shareable()
+            s.set_peer_props({ReservedKey.IDENTITY_NAME: client_name})
+            s.set_header(AppConstants.CONTRIBUTION_ROUND, 0)
+            agg.accept(dxo_collection.update_shareable(s), fl_ctx)
+
+        result = agg.aggregate(fl_ctx)
+        result_dxo = from_shareable(result)
+        for dxo_name in dxo_names:
+            np.testing.assert_allclose(
+                result_dxo.data[dxo_name].data["var1"], weighted_sum[dxo_name] / sum_of_weights[dxo_name]
+            )