From 904e2e9c0c348043a5ad028c3b151df382a14956 Mon Sep 17 00:00:00 2001
From: protolambda <proto@protolambda.com>
Date: Mon, 13 May 2019 23:15:02 +0200
Subject: [PATCH] BLS on/off deco

---
 scripts/phase0/build_spec.py                |  2 +-
 test_libs/pyspec/eth2spec/test/context.py   | 27 ++++++++++++++++++++-
 test_libs/pyspec/eth2spec/test/utils.py     |  4 +++
 test_libs/pyspec/eth2spec/utils/bls.py      | 25 +++++++++++++++++++
 test_libs/pyspec/eth2spec/utils/bls_stub.py | 12 ---------
 5 files changed, 56 insertions(+), 14 deletions(-)
 create mode 100644 test_libs/pyspec/eth2spec/utils/bls.py
 delete mode 100644 test_libs/pyspec/eth2spec/utils/bls_stub.py

diff --git a/scripts/phase0/build_spec.py b/scripts/phase0/build_spec.py
index da5845951..26b0e5a8a 100644
--- a/scripts/phase0/build_spec.py
+++ b/scripts/phase0/build_spec.py
@@ -13,7 +13,7 @@ from typing import (
     Tuple,
 )
 from eth2spec.utils.minimal_ssz import *
-from eth2spec.utils.bls_stub import *
+from eth2spec.utils.bls import *
 
 """)
     for i in (1, 2, 3, 4, 8, 32, 48, 96):
diff --git a/test_libs/pyspec/eth2spec/test/context.py b/test_libs/pyspec/eth2spec/test/context.py
index afabd4a57..3ef62eae5 100644
--- a/test_libs/pyspec/eth2spec/test/context.py
+++ b/test_libs/pyspec/eth2spec/test/context.py
@@ -1,4 +1,5 @@
 from eth2spec.phase0 import spec
+from eth2spec.utils import bls
 
 from .helpers import create_genesis_state
 
@@ -10,7 +11,7 @@ with_state = with_args(lambda: [create_genesis_state(spec.SLOTS_PER_EPOCH * 8, l
 
 # shorthand for decorating @with_state @spectest()
 def spec_state_test(fn):
-    return with_state(spectest()(fn))
+    return with_state(bls_switch(spectest()(fn)))
 
 
 def expect_assertion_error(fn):
@@ -25,3 +26,27 @@ def expect_assertion_error(fn):
         pass
     if bad:
         raise AssertionError('expected an assertion error, but got none.')
+
+
+def always_bls(fn):
+    """
+    Decorator to apply on ``bls_switch`` decorator to force BLS activation. Useful to mark tests as BLS-dependent.
+    """
+    def entry(*args, **kw):
+        # override bls setting
+        kw['bls_active'] = True
+        fn(*args, **kw)
+    return entry
+
+
+def bls_switch(fn):
+    """
+    Decorator to make a function execute with BLS ON, or BLS off.
+    Based on an optional bool argument ``bls_active``, passed to the function at runtime.
+    """
+    def entry(*args, **kw):
+        old_state = bls.bls_active
+        bls.bls_active = kw.pop('bls_active', False)
+        fn(*args, **kw)
+        bls.bls_active = old_state
+    return entry
diff --git a/test_libs/pyspec/eth2spec/test/utils.py b/test_libs/pyspec/eth2spec/test/utils.py
index b19d4df59..1c157bcee 100644
--- a/test_libs/pyspec/eth2spec/test/utils.py
+++ b/test_libs/pyspec/eth2spec/test/utils.py
@@ -1,4 +1,5 @@
 from eth2spec.debug.encode import encode
+from eth2spec.utils import bls
 
 
 def spectest(description: str = None):
@@ -18,6 +19,9 @@ def spectest(description: str = None):
                 else:
                     # description can be explicit
                     out['description'] = description
+                # If BLS is not active, mark the test as BLS-ignorant.
+                if not bls.bls_active:
+                    out['stub_bls'] = True
                 # put all generated data into a dict.
                 for data in fn(*args, **kw):
                     # If there is a type argument, encode it as that type.
diff --git a/test_libs/pyspec/eth2spec/utils/bls.py b/test_libs/pyspec/eth2spec/utils/bls.py
new file mode 100644
index 000000000..daeb361a9
--- /dev/null
+++ b/test_libs/pyspec/eth2spec/utils/bls.py
@@ -0,0 +1,25 @@
+from py_ecc import bls
+
+# Flag to make BLS active or not. Used for testing, do not ignore BLS in production unless you know what you are doing.
+bls_active = True
+
+
+def bls_verify(pubkey, message_hash, signature, domain):
+    if bls_active:
+        return bls.verify(message_hash=message_hash, pubkey=pubkey, signature=signature, domain=domain)
+    else:
+        return True
+
+
+def bls_verify_multiple(pubkeys, message_hashes, signature, domain):
+    if bls_active:
+        return bls.verify_multiple(pubkeys, message_hashes, signature, domain)
+    else:
+        return True
+
+
+def bls_aggregate_pubkeys(pubkeys):
+    if bls_active:
+        return bls.aggregate_pubkeys(pubkeys)
+    else:
+        return b'\x42' * 96
diff --git a/test_libs/pyspec/eth2spec/utils/bls_stub.py b/test_libs/pyspec/eth2spec/utils/bls_stub.py
deleted file mode 100644
index 108c4ef71..000000000
--- a/test_libs/pyspec/eth2spec/utils/bls_stub.py
+++ /dev/null
@@ -1,12 +0,0 @@
-
-
-def bls_verify(pubkey, message_hash, signature, domain):
-    return True
-
-
-def bls_verify_multiple(pubkeys, message_hashes, signature, domain):
-    return True
-
-
-def bls_aggregate_pubkeys(pubkeys):
-    return b'\x42' * 96