From dbf7fbd3d041717f606052424d06e60b933f9843 Mon Sep 17 00:00:00 2001
From: protolambda <proto@protolambda.com>
Date: Thu, 24 Jun 2021 17:13:36 +0200
Subject: [PATCH] encode, decode and randomize ssz Union types

---
 tests/core/pyspec/eth2spec/debug/decode.py     | 13 ++++++++++++-
 tests/core/pyspec/eth2spec/debug/encode.py     |  8 +++++++-
 .../core/pyspec/eth2spec/debug/random_value.py | 18 +++++++++++++++++-
 3 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/tests/core/pyspec/eth2spec/debug/decode.py b/tests/core/pyspec/eth2spec/debug/decode.py
index 871748d54..30bfd487b 100644
--- a/tests/core/pyspec/eth2spec/debug/decode.py
+++ b/tests/core/pyspec/eth2spec/debug/decode.py
@@ -2,7 +2,7 @@ from typing import Any
 from eth2spec.utils.ssz.ssz_impl import hash_tree_root
 from eth2spec.utils.ssz.ssz_typing import (
     uint, Container, List, boolean,
-    Vector, ByteVector, ByteList
+    Vector, ByteVector, ByteList, Union, View
 )
 
 
@@ -27,5 +27,16 @@ def decode(data: Any, typ):
             assert (data["hash_tree_root"][2:] ==
                     hash_tree_root(ret).hex())
         return ret
+    elif issubclass(typ, Union):
+        selector = int(data["selector"])
+        options = typ.options()
+        value_typ = options[selector]
+        value: View
+        if value_typ is None:  # handle the "nil" type case
+            assert data["value"] is None
+            value = None
+        else:
+            value = decode(data["value"], value_typ)
+        return typ(selector=selector, value=value)
     else:
         raise Exception(f"Type not recognized: data={data}, typ={typ}")
diff --git a/tests/core/pyspec/eth2spec/debug/encode.py b/tests/core/pyspec/eth2spec/debug/encode.py
index 10bf4bedd..d93f7cf5e 100644
--- a/tests/core/pyspec/eth2spec/debug/encode.py
+++ b/tests/core/pyspec/eth2spec/debug/encode.py
@@ -1,7 +1,7 @@
 from eth2spec.utils.ssz.ssz_impl import hash_tree_root, serialize
 from eth2spec.utils.ssz.ssz_typing import (
     uint, boolean,
-    Bitlist, Bitvector, Container, Vector, List
+    Bitlist, Bitvector, Container, Vector, List, Union
 )
 
 
@@ -31,5 +31,11 @@ def encode(value, include_hash_tree_roots=False):
         if include_hash_tree_roots:
             ret["hash_tree_root"] = '0x' + hash_tree_root(value).hex()
         return ret
+    elif isinstance(value, Union):
+        inner_value = value.value()
+        return {
+            'selector': int(value.selector()),
+            'value': None if inner_value is None else encode(inner_value, include_hash_tree_roots)
+        }
     else:
         raise Exception(f"Type not recognized: value={value}, typ={type(value)}")
diff --git a/tests/core/pyspec/eth2spec/debug/random_value.py b/tests/core/pyspec/eth2spec/debug/random_value.py
index 33eb25c0c..f8cb5701e 100644
--- a/tests/core/pyspec/eth2spec/debug/random_value.py
+++ b/tests/core/pyspec/eth2spec/debug/random_value.py
@@ -5,7 +5,7 @@ from typing import Type
 
 from eth2spec.utils.ssz.ssz_typing import (
     View, BasicView, uint, Container, List, boolean,
-    Vector, ByteVector, ByteList, Bitlist, Bitvector
+    Vector, ByteVector, ByteList, Bitlist, Bitvector, Union
 )
 
 # in bytes
@@ -115,6 +115,22 @@ def get_random_ssz_object(rng: Random,
                 get_random_ssz_object(rng, field_type, max_bytes_length, max_list_length, mode, chaos)
             for field_name, field_type in fields.items()
         })
+    elif issubclass(typ, Union):
+        options = typ.options()
+        selector: int
+        if mode == RandomizationMode.mode_zero:
+            selector = 0
+        elif mode == RandomizationMode.mode_max:
+            selector = len(options)-1
+        else:
+            selector = rng.randrange(0, len(options))
+        elem_type = options[selector]
+        elem: View
+        if elem_type is None:
+            elem = None
+        else:
+            elem = get_random_ssz_object(rng, elem_type, max_bytes_length, max_list_length, mode, chaos)
+        return typ(selector=selector, value=elem)
     else:
         raise Exception(f"Type not recognized: typ={typ}")