From 2afe8104e0aed72fafa3bf2c9fe7223fe6735834 Mon Sep 17 00:00:00 2001
From: Shawn <44221603+shaspitz@users.noreply.github.com>
Date: Wed, 24 Apr 2024 01:36:12 -0700
Subject: [PATCH] feat: updated example script w/ new queries

---
 .../scripts/ValidatorExampleScript.s.sol      | 129 +++++++++++-------
 1 file changed, 76 insertions(+), 53 deletions(-)

diff --git a/contracts/scripts/ValidatorExampleScript.s.sol b/contracts/scripts/ValidatorExampleScript.s.sol
index ce701bd79..5408d6318 100644
--- a/contracts/scripts/ValidatorExampleScript.s.sol
+++ b/contracts/scripts/ValidatorExampleScript.s.sol
@@ -5,89 +5,112 @@ import "forge-std/Script.sol";
 import "forge-std/console.sol";
 import "../contracts/ValidatorRegistry.sol";
 
-contract StakeAndUnstakeExample is Script {
-    ValidatorRegistry private _validatorRegistry = ValidatorRegistry(0xe7f1725E7734CE288F8367e1Bb143E90bb3F0512);
+abstract contract ExampleScript is Script {
+
+    ValidatorRegistry internal _validatorRegistry = ValidatorRegistry(0xe7f1725E7734CE288F8367e1Bb143E90bb3F0512);
+    address public defaultEOA = 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266;
+
+    function checkStaking(bytes[] memory blsKeys) public view {
+
+        console.log("--------------------");
+        console.log("Checking Staking related state...");
+        console.log("--------------------");
+        
+        for (uint i = 0; i < blsKeys.length; i++) {
+            bool isStaked = _validatorRegistry.isStaked(blsKeys[i]);
+            console.log("--------------------");
+            console.log("BLS Key: ");
+            console.logBytes(blsKeys[i]);
+            console.log("is Staked:", isStaked);
+            uint256 stakedAmount = _validatorRegistry.getStakedAmount(blsKeys[i]);
+            console.log("Staked Amount:", stakedAmount);
+            uint256 unstakingAmount = _validatorRegistry.getUnstakingAmount(blsKeys[i]);
+            console.log("Unstaking Amount:", unstakingAmount);
+            console.log("--------------------");
+        }
+
+        uint256 numStakedValidators = _validatorRegistry.getNumberOfStakedValidators();
+        console.log("Num Staked Validators:", numStakedValidators);
+        if (numStakedValidators == 0) {
+            return;
+        }
+        bytes[] memory vals = _validatorRegistry.getStakedValidators(0, numStakedValidators);
+        for (uint i = 0; i < vals.length; i++) {
+            console.log("Staked validator from batch query: ");
+            console.logBytes(vals[i]);
+        }
+    }
+}
+
+contract StakeExample is ExampleScript {
 
     function run() external {
         vm.startBroadcast();
 
-        console.log("Balance of 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266:", address(0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266).balance);
+        console.log("Balance of 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266:", defaultEOA.balance);
 
         bytes[] memory validators = new bytes[](3);
         validators[0] = hex"a97794deb52ea4529d37d283213ca7e298ea9be0a2fec1bb3134a1464ab8cf9eb2c703d1b42dd68d97b5f1c8e74cc0df";
         validators[1] = hex"96db1884af7bf7a1b57c77222723286a8ce3ef9a16ab6c5542ec5160662d450a1b396b22fc519679adae6ad741547268";
         validators[2] = hex"a840634f574c20e9a35ff80be19309dfd3ace623a093f114e7f44555e3035725b4d1d59b3ce0b2f169871fbe7abc448a";
 
+        checkStaking(validators);
+
         uint256 totalAmount = 10 ether;
         _validatorRegistry.stake{value: totalAmount}(validators);
+        console.log("Stake completed ", totalAmount, "ETH");
 
-        for (uint i = 0; i < validators.length; i++) {
-            checkIsStaked(validators[i]);
-        }
+        checkStaking(validators);
 
-        _validatorRegistry.unstake(validators);
+        console.log("Balance of 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266:", defaultEOA.balance);
+        vm.stopBroadcast();
+    }
+}
 
-        for (uint i = 0; i < validators.length; i++) {
-            checkIsStaked(validators[i]);
-        }
+contract UnstakeExample is ExampleScript {
 
-        console.log("Balance of 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266:", address(0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266).balance);
+    function run() external {
+        vm.startBroadcast();
 
-        vm.stopBroadcast();
-    }
+        console.log("Balance of 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266:", defaultEOA.balance);
 
-    function checkIsStaked(bytes memory blsKey) public view {
-        bool isStaked = _validatorRegistry.isStaked(blsKey);
-        console.log("--------------------");
-        console.log("BLS Key: ");
-        console.logBytes(blsKey);
-        console.log("is Staked:", isStaked);
-        uint256 stakedAmount = _validatorRegistry.getStakedAmount(blsKey);
-        console.log("Staked Amount:", stakedAmount);
-        console.log("--------------------");
+        bytes[] memory validators = new bytes[](3);
+        validators[0] = hex"a97794deb52ea4529d37d283213ca7e298ea9be0a2fec1bb3134a1464ab8cf9eb2c703d1b42dd68d97b5f1c8e74cc0df";
+        validators[1] = hex"96db1884af7bf7a1b57c77222723286a8ce3ef9a16ab6c5542ec5160662d450a1b396b22fc519679adae6ad741547268";
+        validators[2] = hex"a840634f574c20e9a35ff80be19309dfd3ace623a093f114e7f44555e3035725b4d1d59b3ce0b2f169871fbe7abc448a";
+
+        checkStaking(validators);
+
+        _validatorRegistry.unstake(validators);
+        console.log("Unstake initiated");
+
+        checkStaking(validators);
+
+        console.log("Balance of 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266:", defaultEOA.balance);
+        vm.stopBroadcast();
     }
 }
 
-contract WithdrawExample is Script {
-    ValidatorRegistry private _validatorRegistry = ValidatorRegistry(0xe7f1725E7734CE288F8367e1Bb143E90bb3F0512);
-
-    bytes public validator1Key = hex"a97794deb52ea4529d37d283213ca7e298ea9be0a2fec1bb3134a1464ab8cf9eb2c703d1b42dd68d97b5f1c8e74cc0df";
-    bytes public validator2Key = hex"96db1884af7bf7a1b57c77222723286a8ce3ef9a16ab6c5542ec5160662d450a1b396b22fc519679adae6ad741547268";
-    bytes public validator3Key = hex"a840634f574c20e9a35ff80be19309dfd3ace623a093f114e7f44555e3035725b4d1d59b3ce0b2f169871fbe7abc448a";
+contract WithdrawExample is ExampleScript {
 
     function run() external {
         vm.startBroadcast();
 
-        console.log("Balance of 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266:", address(0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266).balance);
-
-        bytes[] memory blsKeys = new bytes[](3);
-        blsKeys[0] = validator1Key;
-        blsKeys[1] = validator2Key;
-        blsKeys[2] = validator3Key;
+        console.log("Balance of 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266:", defaultEOA.balance);
 
-        for (uint i = 0; i < blsKeys.length; i++) {
-            checkIsStaked(blsKeys[i]);
-        }
+        bytes[] memory validators = new bytes[](3);
+        validators[0] = hex"a97794deb52ea4529d37d283213ca7e298ea9be0a2fec1bb3134a1464ab8cf9eb2c703d1b42dd68d97b5f1c8e74cc0df";
+        validators[1] = hex"96db1884af7bf7a1b57c77222723286a8ce3ef9a16ab6c5542ec5160662d450a1b396b22fc519679adae6ad741547268";
+        validators[2] = hex"a840634f574c20e9a35ff80be19309dfd3ace623a093f114e7f44555e3035725b4d1d59b3ce0b2f169871fbe7abc448a";
 
-        _validatorRegistry.withdraw(blsKeys);
+        checkStaking(validators);
 
-        for (uint i = 0; i < blsKeys.length; i++) {
-            checkIsStaked(blsKeys[i]);
-        }
+        _validatorRegistry.withdraw(validators);
+        console.log("Withdraw initiated");
 
-        console.log("Balance of 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266:", address(0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266).balance);
+        checkStaking(validators);
 
+        console.log("Balance of 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266:", defaultEOA.balance);
         vm.stopBroadcast();
     }
-
-    function checkIsStaked(bytes memory blsKey) public view {
-        bool isStaked = _validatorRegistry.isStaked(blsKey);
-        console.log("--------------------");
-        console.log("BLS Key: ");
-        console.logBytes(blsKey);
-        console.log("is Staked:", isStaked);
-        uint256 stakedAmount = _validatorRegistry.getStakedAmount(blsKey);
-        console.log("Staked Amount:", stakedAmount);
-        console.log("--------------------");
-    }
 }