From 8e303b60b961420bb1df6b8ea93be5521a32e02a Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hsiaowei.eth@gmail.com>
Date: Fri, 12 Mar 2021 00:06:52 +0800
Subject: [PATCH] Fix generator to ssz_snappy mode

---
 tests/generators/fork_choice/main.py          | 45 +++++++------------
 tests/generators/fork_choice/requirements.txt |  4 +-
 2 files changed, 17 insertions(+), 32 deletions(-)

diff --git a/tests/generators/fork_choice/main.py b/tests/generators/fork_choice/main.py
index f684399eb..79920d8ad 100644
--- a/tests/generators/fork_choice/main.py
+++ b/tests/generators/fork_choice/main.py
@@ -1,40 +1,25 @@
-from typing import Iterable
-
-from gen_base import gen_runner, gen_typing
-from gen_from_tests.gen import generate_from_tests
-from importlib import reload, import_module
-from eth2spec.config import config_util
+from eth2spec.gen_helpers.gen_from_tests.gen import run_state_test_generators
 from eth2spec.phase0 import spec as spec_phase0
-from eth2spec.test.context import PHASE0
-from eth2spec.utils import bls
+from eth2spec.lightclient_patch import spec as spec_lightclient_patch
+from eth2spec.phase1 import spec as spec_phase1
+from eth2spec.test.context import PHASE0, PHASE1, LIGHTCLIENT_PATCH
 
 
-def create_provider(fork_name: str, handler_name: str,
-                    tests_src_mod_name: str, config_name: str) -> gen_typing.TestProvider:
-    def prepare_fn(configs_path: str) -> str:
-        config_util.prepare_config(configs_path, config_name)
-        reload(spec_phase0)
-        bls.use_milagro()
-        return config_name
-
-    def cases_fn() -> Iterable[gen_typing.TestCase]:
-        tests_src = import_module(tests_src_mod_name)
-        return generate_from_tests(
-            runner_name='fork_choice',
-            handler_name=handler_name,
-            src=tests_src,
-            fork_name=fork_name,
-        )
-
-    return gen_typing.TestProvider(prepare=prepare_fn, make_cases=cases_fn)
+specs = (spec_phase0, spec_lightclient_patch, spec_phase1)
 
 
 if __name__ == "__main__":
     phase_0_mods = {key: 'eth2spec.test.phase0.fork_choice.test_' + key for key in [
         'get_head',
     ]}
+    # No additional lightclient_patch or phase 1 specific finality tests, yet.
+    lightclient_patch_mods = phase_0_mods
+    phase_1_mods = phase_0_mods
 
-    # TODO: add other configs and forks
-    gen_runner.run_generator(f"fork_choice", [
-        create_provider(PHASE0, key, mod_name, 'minimal') for key, mod_name in phase_0_mods.items()
-    ])
+    all_mods = {
+        PHASE0: phase_0_mods,
+        LIGHTCLIENT_PATCH: lightclient_patch_mods,
+        PHASE1: phase_1_mods,
+    }
+
+    run_state_test_generators(runner_name="fork_choice", specs=specs, all_mods=all_mods)
diff --git a/tests/generators/fork_choice/requirements.txt b/tests/generators/fork_choice/requirements.txt
index b82314298..182248686 100644
--- a/tests/generators/fork_choice/requirements.txt
+++ b/tests/generators/fork_choice/requirements.txt
@@ -1,2 +1,2 @@
-../../core/gen_helpers
-../../../
\ No newline at end of file
+pytest>=4.4
+../../../[generator]