From 6b1d4839d846c786ab0e864165aab1dbe5ee3b67 Mon Sep 17 00:00:00 2001
From: protolambda <proto@protolambda.com>
Date: Tue, 7 Jul 2020 20:17:36 +0200
Subject: [PATCH] update block operations generator

---
 tests/generators/operations/main.py | 61 +++++++++++++++++------------
 1 file changed, 35 insertions(+), 26 deletions(-)

diff --git a/tests/generators/operations/main.py b/tests/generators/operations/main.py
index 6d4f6d139..b8493a21b 100644
--- a/tests/generators/operations/main.py
+++ b/tests/generators/operations/main.py
@@ -1,25 +1,16 @@
 from typing import Iterable
 
-from eth2spec.test.phase0.block_processing import (
-    test_process_attestation,
-    test_process_attester_slashing,
-    test_process_block_header,
-    test_process_deposit,
-    test_process_proposer_slashing,
-    test_process_voluntary_exit,
-)
-
+from importlib import import_module
 from gen_base import gen_runner, gen_typing
 from gen_from_tests.gen import generate_from_tests
 from importlib import reload
 from eth2spec.config import config_util
 from eth2spec.phase0 import spec as spec_phase0
 from eth2spec.phase1 import spec as spec_phase1
-from eth2spec.test.context import PHASE0
+from eth2spec.test.context import PHASE0, PHASE1
 
 
-def create_provider(handler_name: str, tests_src, config_name: str) -> gen_typing.TestProvider:
-
+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)
@@ -27,28 +18,46 @@ def create_provider(handler_name: str, tests_src, config_name: str) -> gen_typin
         return config_name
 
     def cases_fn() -> Iterable[gen_typing.TestCase]:
+        tests_src = import_module(tests_src_mod_name)
         return generate_from_tests(
             runner_name='operations',
             handler_name=handler_name,
             src=tests_src,
-            fork_name=PHASE0,
+            fork_name=fork_name,
         )
 
     return gen_typing.TestProvider(prepare=prepare_fn, make_cases=cases_fn)
 
 
 if __name__ == "__main__":
-    gen_runner.run_generator("operations", [
-        create_provider('attestation', test_process_attestation, 'minimal'),
-        create_provider('attestation', test_process_attestation, 'mainnet'),
-        create_provider('attester_slashing', test_process_attester_slashing, 'minimal'),
-        create_provider('attester_slashing', test_process_attester_slashing, 'mainnet'),
-        create_provider('block_header', test_process_block_header, 'minimal'),
-        create_provider('block_header', test_process_block_header, 'mainnet'),
-        create_provider('deposit', test_process_deposit, 'minimal'),
-        create_provider('deposit', test_process_deposit, 'mainnet'),
-        create_provider('proposer_slashing', test_process_proposer_slashing, 'minimal'),
-        create_provider('proposer_slashing', test_process_proposer_slashing, 'mainnet'),
-        create_provider('voluntary_exit', test_process_voluntary_exit, 'minimal'),
-        create_provider('voluntary_exit', test_process_voluntary_exit, 'mainnet'),
+    phase_0_mods = [(key, 'eth2spec.test.phase0.block_processing.test_process_' + key) for key in [
+        'attestation',
+        'attester_slashing',
+        'block_header',
+        'deposit',
+        'proposer_slashing',
+        'voluntary_exit',
+    ]]
+    phase_1_mods = [(key, 'eth2spec.test.phase1.block_processing.test_process_' + key) for key in [
+        'attestation',
+        'chunk_challenge',
+        'custody_key_reveal',
+        'custody_slashing',
+        'early_derived_secret_reveal',
+        'shard_transition',
+    ]] + phase_0_mods  # also run the previous phase 0 tests (but against phase 1 spec)
+
+    gen_runner.run_generator(f"phase0/operations", [
+        create_provider(PHASE0, key, mod_name, 'minimal') for key, mod_name in phase_0_mods
     ])
+    # TODO: disabled for testing
+    # gen_runner.run_generator(f"phase0/operations", [
+    #     create_provider(key, mod_name, 'mainnet') for key, mod_name in phase_0_mods
+    # ])
+    gen_runner.run_generator(f"phase1/operations", [
+        create_provider(PHASE1, key, mod_name, 'minimal') for key, mod_name in phase_1_mods
+    ])
+    # Disabled for now
+    # gen_runner.run_generator(f"phase1/operations", [
+    #     create_provider(PHASE1, key, mod_name, 'mainnet') for key, mod_name in phase_1_mods
+    # ])