From 2477deaf1312539cc0e429452ffe3e2d608e6b28 Mon Sep 17 00:00:00 2001
From: Alex Stokes <r.alex.stokes@gmail.com>
Date: Mon, 30 Aug 2021 13:53:05 -0700
Subject: [PATCH] Allow test driver to batch test cases under one handler name
 with a list

---
 .../gen_helpers/gen_from_tests/gen.py         | 25 +++++++++++--------
 tests/generators/operations/main.py           |  6 ++---
 2 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/tests/core/pyspec/eth2spec/gen_helpers/gen_from_tests/gen.py b/tests/core/pyspec/eth2spec/gen_helpers/gen_from_tests/gen.py
index db3f9e949..88bc6d601 100644
--- a/tests/core/pyspec/eth2spec/gen_helpers/gen_from_tests/gen.py
+++ b/tests/core/pyspec/eth2spec/gen_helpers/gen_from_tests/gen.py
@@ -1,6 +1,6 @@
 from importlib import import_module
 from inspect import getmembers, isfunction
-from typing import Any, Callable, Dict, Iterable, Optional
+from typing import Any, Callable, Dict, Iterable, Optional, List, Union
 
 from eth2spec.utils import bls
 from eth2spec.test.helpers.constants import ALL_PRESETS, TESTGEN_FORKS
@@ -59,8 +59,10 @@ def generate_from_tests(runner_name: str, handler_name: str, src: Any,
 def get_provider(create_provider_fn: Callable[[SpecForkName, PresetBaseName, str, str], TestProvider],
                  fork_name: SpecForkName,
                  preset_name: PresetBaseName,
-                 all_mods: Dict[str, Dict[str, str]]) -> Iterable[TestProvider]:
+                 all_mods: Dict[str, Dict[str, Union[List[str], str]]]) -> Iterable[TestProvider]:
     for key, mod_name in all_mods[fork_name].items():
+        if not isinstance(mod_name, List):
+            mod_name = [mod_name]
         yield create_provider_fn(
             fork_name=fork_name,
             preset_name=preset_name,
@@ -75,16 +77,17 @@ def get_create_provider_fn(runner_name: str) -> Callable[[SpecForkName, str, str
         return
 
     def create_provider(fork_name: SpecForkName, preset_name: PresetBaseName,
-                        handler_name: str, tests_src_mod_name: str) -> TestProvider:
+                        handler_name: str, tests_src_mod_name: List[str]) -> TestProvider:
         def cases_fn() -> Iterable[TestCase]:
-            tests_src = import_module(tests_src_mod_name)
-            return generate_from_tests(
-                runner_name=runner_name,
-                handler_name=handler_name,
-                src=tests_src,
-                fork_name=fork_name,
-                preset_name=preset_name,
-            )
+            for mod_name in tests_src_mod_name:
+                tests_src = import_module(mod_name)
+                yield from generate_from_tests(
+                    runner_name=runner_name,
+                    handler_name=handler_name,
+                    src=tests_src,
+                    fork_name=fork_name,
+                    preset_name=preset_name,
+                )
 
         return TestProvider(prepare=prepare_fn, make_cases=cases_fn)
     return create_provider
diff --git a/tests/generators/operations/main.py b/tests/generators/operations/main.py
index d2653d87d..0f4da69a5 100644
--- a/tests/generators/operations/main.py
+++ b/tests/generators/operations/main.py
@@ -12,9 +12,9 @@ if __name__ == "__main__":
         'voluntary_exit',
     ]}
     altair_mods = {
-        **{key: 'eth2spec.test.altair.block_processing.sync_aggregate.test_process_' + key for key in [
-            'sync_aggregate',
-            'sync_aggregate_random',
+        **{'sync_aggregate': [
+            'eth2spec.test.altair.block_processing.sync_aggregate.test_process_' + key
+            for key in ['sync_aggregate', 'sync_aggregate_random']
         ]},
         **phase_0_mods,
     }  # also run the previous phase 0 tests