Remote-db: fix keys per validator calculation (stakewise#241)

* Remote-db: fix keys per validator calculation Signed-off-by: cyc60 <[email protected]> * Review fix Signed-off-by: cyc60 <[email protected]> --------- Signed-off-by: cyc60 <[email protected]>
ChorusOne · Nov 23, 2023 · f7fd19e · f7fd19e
1 parent 30cb5a9
commit f7fd19e
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 4 deletions.
diff --git a/src/remote_db/tasks.py b/src/remote_db/tasks.py
@@ -195,9 +195,11 @@ def setup_validator(
         raise click.ClickException('No keypairs found in the remote db.')
 
     public_keys_count = len(keypairs)
-    keys_per_validator = public_keys_count // total_validators
-    start_index = keys_per_validator * validator_index
-    end_index = min(start_index + keys_per_validator, public_keys_count)
+    start_index, end_index = _get_key_indexes(
+        public_keys_count=public_keys_count,
+        validator_index=validator_index,
+        total_validators=total_validators,
+    )
     if not 0 <= start_index < end_index <= public_keys_count:
         raise click.ClickException('Invalid validator index')
 
@@ -348,3 +350,14 @@ def _generate_proposer_config(
     }
     with open(filepath, 'w', encoding='utf-8') as f:
         json.dump(config, f, ensure_ascii=False, indent=4)
+
+
+def _get_key_indexes(
+    public_keys_count: int,
+    total_validators: int,
+    validator_index: int,
+) -> tuple[int, int]:
+    keys_per_validator = (public_keys_count - 1) // total_validators + 1
+    start_index = keys_per_validator * validator_index
+    end_index = min(start_index + keys_per_validator, public_keys_count)
+    return start_index, end_index
diff --git a/src/remote_db/tests/test_tasks.py b/src/remote_db/tests/test_tasks.py
@@ -15,7 +15,7 @@
 from src.common.typings import Oracles
 from src.remote_db.commands import remote_db_group
 from src.remote_db.database import ConfigsCrud, KeyPairsCrud
-from src.remote_db.tasks import _encrypt_private_key
+from src.remote_db.tasks import _encrypt_private_key, _get_key_indexes
 from src.remote_db.typings import RemoteDatabaseKeyPair
 from src.validators.signing.key_shares import private_key_to_private_key_shares
 from src.validators.typings import BLSPrivkey
@@ -395,3 +395,29 @@ def test_setup_operator(
             result = runner.invoke(remote_db_group, args)
             output = 'Successfully created operator configuration file.\n'
             assert output.strip() in result.output.strip()
+
+
+def test_get_key_indexes():
+    assert _get_key_indexes(1, 1, 0) == (0, 1)
+
+    assert _get_key_indexes(2, 1, 0) == (0, 2)
+
+    assert _get_key_indexes(2, 2, 0) == (0, 1)
+    assert _get_key_indexes(2, 2, 1) == (1, 2)
+
+    assert _get_key_indexes(27, 2, 0) == (0, 14)
+    assert _get_key_indexes(27, 2, 1) == (14, 27)
+
+    assert _get_key_indexes(150, 3, 0) == (0, 50)
+    assert _get_key_indexes(150, 3, 1) == (50, 100)
+    assert _get_key_indexes(150, 3, 2) == (100, 150)
+
+    total = 50
+    count = 0
+    for i in range(total):
+        if i == 49:
+            assert _get_key_indexes(199, total, i) == (196, 199)
+        else:
+            assert _get_key_indexes(199, total, i) == (i * 4, i * 4 + 4)
+        count += _get_key_indexes(199, total, i)[1] - _get_key_indexes(199, total, i)[0]
+    assert count == 199