From 745d529598632029fee9820590b033b7e0e23935 Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hsiaowei.eth@gmail.com>
Date: Tue, 25 Apr 2023 12:57:42 +0800
Subject: [PATCH] Add `compute_subscribed_subnets` unittests and fix typing
 errors

---
 specs/phase0/validator.md                     | 10 +++--
 .../validator/test_validator_unittest.py      | 39 ++++++++++++++++++-
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/specs/phase0/validator.md b/specs/phase0/validator.md
index 56ca50732..5266fec7a 100644
--- a/specs/phase0/validator.md
+++ b/specs/phase0/validator.md
@@ -616,10 +616,14 @@ Because Phase 0 does not have shards and thus does not have Shard Committees, th
 
 ```python
 def compute_subscribed_subnet(node_id: int, epoch: Epoch, index: int) -> int:
-    node_id_prefix = node_id >> (256 - ATTESTATION_SUBNET_PREFIX_BITS)
+    node_id_prefix = node_id >> (256 - int(ATTESTATION_SUBNET_PREFIX_BITS))
     node_offset = node_id % EPOCHS_PER_SUBNET_SUBSCRIPTION
-    permutation_seed = hash(uint_to_bytes((epoch + node_offset) // EPOCHS_PER_SUBNET_SUBSCRIPTION))
-    permutated_prefix = compute_shuffled_index(node_id_prefix, 1 << ATTESTATION_SUBNET_PREFIX_BITS, permutation_seed)
+    permutation_seed = hash(uint_to_bytes(uint64((epoch + node_offset) // EPOCHS_PER_SUBNET_SUBSCRIPTION)))
+    permutated_prefix = compute_shuffled_index(
+        node_id_prefix,
+        1 << int(ATTESTATION_SUBNET_PREFIX_BITS),
+        permutation_seed,
+    )
     return (permutated_prefix + index) % ATTESTATION_SUBNET_COUNT
 ```
 
diff --git a/tests/core/pyspec/eth2spec/test/phase0/unittests/validator/test_validator_unittest.py b/tests/core/pyspec/eth2spec/test/phase0/unittests/validator/test_validator_unittest.py
index cf7ef392f..177748eac 100644
--- a/tests/core/pyspec/eth2spec/test/phase0/unittests/validator/test_validator_unittest.py
+++ b/tests/core/pyspec/eth2spec/test/phase0/unittests/validator/test_validator_unittest.py
@@ -1,6 +1,12 @@
+import random
+
 from eth2spec.test.context import (
+    single_phase,
     spec_state_test,
-    always_bls, with_phases, with_all_phases,
+    spec_test,
+    always_bls,
+    with_phases,
+    with_all_phases,
 )
 from eth2spec.test.helpers.constants import PHASE0
 from eth2spec.test.helpers.attestations import build_attestation_data, get_valid_attestation
@@ -476,3 +482,34 @@ def test_get_aggregate_and_proof_signature(spec, state):
         privkey=privkey,
         pubkey=pubkey,
     )
+
+
+def run_compute_subscribed_subnets_arguments(spec, rng=random.Random(1111)):
+    node_id = rng.randint(0, 2**40 - 1)  # try VALIDATOR_REGISTRY_LIMIT
+    epoch = rng.randint(0, 2**64 - 1)
+    subnets = spec.compute_subscribed_subnets(node_id, epoch)
+    assert len(subnets) == spec.SUBNETS_PER_NODE
+
+
+@with_all_phases
+@spec_test
+@single_phase
+def test_compute_subscribed_subnets_random_1(spec):
+    rng = random.Random(1111)
+    run_compute_subscribed_subnets_arguments(spec, rng)
+
+
+@with_all_phases
+@spec_test
+@single_phase
+def test_compute_subscribed_subnets_random_2(spec):
+    rng = random.Random(2222)
+    run_compute_subscribed_subnets_arguments(spec, rng)
+
+
+@with_all_phases
+@spec_test
+@single_phase
+def test_compute_subscribed_subnets_random_3(spec):
+    rng = random.Random(3333)
+    run_compute_subscribed_subnets_arguments(spec, rng)