upload csm keys with batches

fix rewards calculation with csm alive
lidofinance · Oct 29, 2024 · e99d8dd · e99d8dd
1 parent f8cd02a
commit e99d8dd
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 8 deletions.
diff --git a/tests/regression/test_pause_resume.py b/tests/regression/test_pause_resume.py
@@ -278,12 +278,14 @@ def test_paused_staking_module_can_reward(burner: Contract, stranger):
     print(report_tx.events["Transfer"])
     module_index = 0
     simple_dvt_index = 1
+    csm_index = 2
 
     if report_tx.events["Transfer"][module_index]["to"] == burner.address:
         module_index += 1
         simple_dvt_index += 1
+        csm_index += 1
 
-    agent_index = module_index + 2
+    agent_index = module_index + 3
     assert report_tx.events["Transfer"][module_index]["to"] == module_address
     assert report_tx.events["Transfer"][module_index]["from"] == ZERO_ADDRESS
     assert report_tx.events["Transfer"][agent_index]["to"] == contracts.agent
@@ -299,8 +301,16 @@ def test_paused_staking_module_can_reward(burner: Contract, stranger):
         * simple_dvt_stats["treasuryFee"]
         // 100_00
     )
+    csm_stats = contracts.staking_router.getStakingModule(3)
+    csm_treasury_fee = (
+        report_tx.events["Transfer"][csm_index]["value"]
+        * 100_00
+        // csm_stats["stakingModuleFee"]
+        * csm_stats["treasuryFee"]
+        // 100_00
+    )
     assert almostEqWithDiff(
-        report_tx.events["Transfer"][module_index]["value"] + simple_dvt_treasury_fee,
+        report_tx.events["Transfer"][module_index]["value"] + simple_dvt_treasury_fee + csm_treasury_fee,
         report_tx.events["Transfer"][agent_index]["value"],
         100,
     )

diff --git a/utils/test/csm_helpers.py b/utils/test/csm_helpers.py
@@ -130,6 +130,7 @@ def get_ea_members():
                 "0xf99292f2fccc3f4fa8cef73512bd5f21a46214d922dfed8e8a65f763c13dfb72"])
     )
 
+
 def csm_add_node_operator(csm, accounting, node_operator, proof, keys_count=5, curve_id=0):
     pubkeys_batch = random_pubkeys_batch(keys_count)
     signatures_batch = random_signatures_batch(keys_count)
@@ -149,14 +150,23 @@ def csm_add_node_operator(csm, accounting, node_operator, proof, keys_count=5, c
 
     return csm.getNodeOperatorsCount() - 1
 
+
 def csm_upload_keys(csm, accounting, no_id, keys_count=5):
     manager_address = csm.getNodeOperator(no_id)["managerAddress"]
-    pubkeys_batch = random_pubkeys_batch(keys_count)
-    signatures_batch = random_signatures_batch(keys_count)
-    csm.addValidatorKeysETH(no_id, keys_count, pubkeys_batch, signatures_batch,
-                            {"from": manager_address,
-                             "value": accounting.getRequiredBondForNextKeys(no_id, keys_count)}
-                            )
+    set_balance_in_wei(manager_address, accounting.getRequiredBondForNextKeys(no_id, keys_count) + ETH(1))
+
+    keys_batch = 100
+    remaining_keys = keys_count
+    while remaining_keys > 0:
+        keys_batch = min(keys_batch, remaining_keys)
+        pubkeys_batch = random_pubkeys_batch(keys_batch)
+        signatures_batch = random_signatures_batch(keys_batch)
+        value = accounting.getRequiredBondForNextKeys(no_id, keys_count)
+        csm.addValidatorKeysETH(no_id, keys_batch, pubkeys_batch, signatures_batch, {
+            "from": csm.getNodeOperator(no_id)["managerAddress"],
+            "value": value
+        })
+        remaining_keys -= keys_batch
 
 
 def fill_csm_operators_with_keys(target_operators_count, keys_count):