From 4d8aeb869926e1d8a71b0230bfd017508ae2a06d Mon Sep 17 00:00:00 2001
From: Alex Vlasov <avv49@mail.ru>
Date: Wed, 3 Apr 2024 02:18:29 +0400
Subject: [PATCH] Added `test_gen.py` driver to generate all tests from a
 single confi. NB: to be renamed to `main.py` in future.

---
 .../filter_block_tree_generator.py            |  7 +-
 .../generators/fork_choice_generated/main.py  |  7 +-
 .../fork_choice_generated/test_gen.py         | 86 +++++++++++++++++++
 3 files changed, 94 insertions(+), 6 deletions(-)
 create mode 100644 tests/generators/fork_choice_generated/test_gen.py

diff --git a/tests/generators/fork_choice_generated/filter_block_tree_generator.py b/tests/generators/fork_choice_generated/filter_block_tree_generator.py
index 7b0c62f48b..a1ffab1815 100644
--- a/tests/generators/fork_choice_generated/filter_block_tree_generator.py
+++ b/tests/generators/fork_choice_generated/filter_block_tree_generator.py
@@ -15,6 +15,10 @@
 GENERATOR_NAME = 'filter_block_tree'
 
 
+forks = [ALTAIR]
+presets = [MINIMAL]
+
+
 def _load_model_solutions(instance_path: str):
     yaml = YAML(typ='safe')
     solutions = yaml.load(open(instance_path, 'r'))
@@ -134,9 +138,6 @@ def make_cases_fn() -> Iterable[TestCase]:
 
 
 if __name__ == "__main__":
-    forks = [ALTAIR]
-    presets = [MINIMAL]
-
     arg_parser = gen_runner.create_arg_parser(GENERATOR_NAME)
 
     arg_parser.add_argument(
diff --git a/tests/generators/fork_choice_generated/main.py b/tests/generators/fork_choice_generated/main.py
index 770586057b..863fffa4d8 100644
--- a/tests/generators/fork_choice_generated/main.py
+++ b/tests/generators/fork_choice_generated/main.py
@@ -18,6 +18,10 @@
 GENERATOR_NAME = 'fork_choice_generated'
 
 
+forks = [ALTAIR]
+presets = [MINIMAL]
+
+
 def _import_test_fn():
     src = import_module('eth2spec.test.phase0.fork_choice.test_sm_links_tree_model')
     print("generating test vectors from tests source: %s" % src.__name__)
@@ -95,9 +99,6 @@ def make_cases_fn() -> Iterable[TestCase]:
 
 
 if __name__ == "__main__":
-    forks = [ALTAIR]
-    presets = [MINIMAL]
-
     arg_parser = gen_runner.create_arg_parser(GENERATOR_NAME)
 
     arg_parser.add_argument(
diff --git a/tests/generators/fork_choice_generated/test_gen.py b/tests/generators/fork_choice_generated/test_gen.py
new file mode 100644
index 0000000000..6815025bf9
--- /dev/null
+++ b/tests/generators/fork_choice_generated/test_gen.py
@@ -0,0 +1,86 @@
+from eth2spec.gen_helpers.gen_base import gen_runner
+
+from filter_block_tree_generator import (
+    forks as block_cover_forks,
+    presets as block_cover_presets,
+    _load_model_solutions as block_cover_load_solutions,
+    _create_providers as block_cover_create_providers
+)
+from main import (
+    forks as block_tree_forks,
+    presets as block_tree_presets,
+    _load_sm_link_solutions as block_tree_load_solutions,
+    _create_providers as block_tree_create_providers
+)
+
+
+GENERATOR_NAME = 'fork_choice_generated'
+
+
+test_gen_config = {
+    'block_tree_test': {
+        'test_type': 'block_tree',
+        'instances': 'block_tree.yaml',
+        'seed': 123,
+        'nr_variations': 3,
+    },
+    'block_cover_test': {
+        'test_type': 'block_cover',
+        'instances': 'block_cover.yaml',
+        'seed': 456,
+        'nr_variations': 3,
+    }
+}
+
+if __name__ == "__main__":
+    arg_parser = gen_runner.create_arg_parser(GENERATOR_NAME)
+
+    arg_parser.add_argument(
+        '--fc-gen-debug',
+        dest='fc_gen_debug',
+        action='store_true',
+        default=False,
+        required=False,
+        help='If set provides debug output and enable additional checks for generated chains',
+    )
+    # TODO: load test_gen_config from a yaml file
+    # arg_parser.add_argument(
+    #     '--fc-gen-config',
+    #     dest='fc_gen_config',
+    #     type=str,
+    #     default=None,
+    #     required=False,
+    #     help='Path to a file with test generator configurations'
+    # )
+
+    args = arg_parser.parse_args()
+
+    for test_name, params in test_gen_config.items():
+        print(test_name)
+        test_type = params['test_type']
+        instances_path = params['instances']
+        initial_seed = params['seed']
+        nr_variations = params['nr_variations']
+
+        if test_type == 'block_tree':
+            solutions = block_tree_load_solutions(instances_path)
+            providers = block_tree_create_providers(forks=block_tree_forks,
+                                                    presets=block_tree_presets,
+                                                    debug=args.fc_gen_debug,
+                                                    initial_seed=initial_seed,
+                                                    solutions=solutions,
+                                                    number_of_variations=nr_variations)
+        elif test_type == 'block_cover':
+            solutions = block_cover_load_solutions(instances_path)
+            providers = block_cover_create_providers(forks=block_cover_forks,
+                                                     presets=block_cover_presets,
+                                                     debug=args.fc_gen_debug,
+                                                     initial_seed=initial_seed,
+                                                     solutions=solutions,
+                                                     number_of_variations=nr_variations)
+        else:
+            raise ValueError(f'Unsupported test type: {test_type}')
+        
+        gen_runner.run_generator(GENERATOR_NAME, providers, arg_parser)
+
+