From ef13f6be3514d22e2d9e0eeee6584e87fc3ac730 Mon Sep 17 00:00:00 2001
From: shr1ftyy <syeamtechdemon@gmail.com>
Date: Sat, 25 Jan 2025 17:16:43 +0600
Subject: [PATCH] feat: tests for get total stake coldkey and hotkey

---
 src/evm/eth.staking.precompile.test.js | 92 ++++++++++++++++++++++++++
 src/util/precompile.js                 | 38 +++++++++++
 2 files changed, 130 insertions(+)

diff --git a/src/evm/eth.staking.precompile.test.js b/src/evm/eth.staking.precompile.test.js
index 3f7f1b1..e86cd91 100644
--- a/src/evm/eth.staking.precompile.test.js
+++ b/src/evm/eth.staking.precompile.test.js
@@ -328,5 +328,97 @@ describe("Staking precompile", () => {
       proxies = await api.query.proxy.proxies(publicKey);
       expect(proxies[0].length).to.be.eq(0);
     });
+
+  });
+
+  it("Can get total stake for coldkey", async () => {
+    await usingEthApi(async (provider) => {
+      const ss58mirror = convertH160ToSS58(fundedEthWallet.address);
+      const netuid = 1;
+
+      let stakeBefore;
+      await usingApi(async (api) => {
+        stakeBefore = u256toBigNumber(
+          await api.query.subtensorModule.alpha(
+            hotkey.address,
+            ss58mirror,
+            netuid
+          )
+        );
+      });
+
+      await usingEthApi(async (provider) => {
+        // Create a contract instances
+        const signer = new ethers.Wallet(fundedEthWallet.privateKey, provider);
+        const contract = new ethers.Contract(
+          ISTAKING_ADDRESS,
+          IStakingABI,
+          signer
+        );
+
+        // Add stake
+        const tx = await contract.addStake(hotkey.publicKey, netuid, {
+          value: amountStr,
+        });
+        await tx.wait();
+
+        await usingApi(async (api) => {
+          const coldPublicKey = convertH160ToPublicKey(fundedEthWallet.address);
+          const coldkeyStake = new BigNumber(
+            await contract.getTotalColdkeyStake(coldPublicKey)
+          );
+
+          // log the coldkey stake
+          console.log("Coldkey stake: ", coldkeyStake.toString());
+
+          expect(coldkeyStake).to.be.bignumber.gt(stakeBefore);
+        });
+      });
+    });
+  });
+
+  it("Can get total stake for hotkey", async () => {
+    await usingEthApi(async (provider) => {
+      const ss58mirror = convertH160ToSS58(fundedEthWallet.address);
+      const netuid = 1;
+
+      let stakeBefore;
+      await usingApi(async (api) => {
+        stakeBefore = u256toBigNumber(
+          await api.query.subtensorModule.alpha(
+            hotkey.address,
+            ss58mirror,
+            netuid
+          )
+        );
+      });
+
+      await usingEthApi(async (provider) => {
+        // Create a contract instances
+        const signer = new ethers.Wallet(fundedEthWallet.privateKey, provider);
+        const contract = new ethers.Contract(
+          ISTAKING_ADDRESS,
+          IStakingABI,
+          signer
+        );
+
+        // Add stake
+        const tx = await contract.addStake(hotkey.publicKey, netuid, {
+          value: amountStr,
+        });
+        await tx.wait();
+
+        await usingApi(async (api) => {
+          const hotkeyStake = new BigNumber(
+            await contract.getTotalHotkeyStake(hotkey.publicKey)
+          );
+
+          // log the hotkey stake
+          console.log("Hotkey stake: ", hotkeyStake.toString());
+
+          expect(hotkeyStake).to.be.bignumber.gt(stakeBefore);
+        });
+      });
+    });
   });
 });
diff --git a/src/util/precompile.js b/src/util/precompile.js
index 0ea4d68..f2c045c 100644
--- a/src/util/precompile.js
+++ b/src/util/precompile.js
@@ -100,6 +100,44 @@ export const IStakingABI = [
     stateMutability: "nonpayable",
     type: "function",
   },
+  {
+    "inputs": [
+      {
+        "internalType": "bytes32",
+        "name": "coldkey",
+        "type": "bytes32"
+      }
+    ],
+    "name": "getTotalColdkeyStake",
+    "outputs": [
+      {
+        "internalType": "uint256",
+        "name": "",
+        "type": "uint256"
+      }
+    ],
+    "stateMutability": "view",
+    "type": "function"
+  },
+  {
+    "inputs": [
+      {
+        "internalType": "bytes32",
+        "name": "hotkey",
+        "type": "bytes32"
+      }
+    ],
+    "name": "getTotalHotkeyStake",
+    "outputs": [
+      {
+        "internalType": "uint256",
+        "name": "",
+        "type": "uint256"
+      }
+    ],
+    "stateMutability": "view",
+    "type": "function"
+  }
 ];
 
 export const IMetagraphABI = [