From 7711bcee04d0177cbacc02d663bb3ca04364de2f Mon Sep 17 00:00:00 2001
From: Ales Tsurko <ales.tsurko@gmail.com>
Date: Tue, 14 Jan 2025 23:32:11 +0100
Subject: [PATCH] Add staking proxy test

---
 package.json                           |   1 +
 src/evm/eth.staking.precompile.test.js | 104 ++++++++++++++++---------
 yarn.lock                              |   5 ++
 3 files changed, 75 insertions(+), 35 deletions(-)

diff --git a/package.json b/package.json
index 7beb44c..25717cb 100644
--- a/package.json
+++ b/package.json
@@ -17,6 +17,7 @@
     "bignumber.js": "^9.1.2",
     "chai": "^5.1.1",
     "chai-as-promised": "^8.0.0",
+    "chai-bignumber": "^3.1.0",
     "ethers": "^6.13.3",
     "mocha": "^10.7.0",
     "ts-node": "^10.9.2",
diff --git a/src/evm/eth.staking.precompile.test.js b/src/evm/eth.staking.precompile.test.js
index 95760b2..74dfd94 100644
--- a/src/evm/eth.staking.precompile.test.js
+++ b/src/evm/eth.staking.precompile.test.js
@@ -7,10 +7,8 @@ import {
   convertH160ToPublicKey,
 } from "../util/address.js";
 import { getExistentialDeposit, getTaoBalance } from "../util/helpers.js";
-import { assert, ethers } from "ethers";
-import BigNumber from "bignumber.js";
+import { ethers } from "ethers";
 import { expect } from "chai";
-import exp from "constants";
 import { getRandomKeypair } from "../util/known-keys.js";
 
 let tk;
@@ -23,67 +21,69 @@ let abi = [
     inputs: [
       {
         internalType: "bytes32",
-        name: "hotkey",
-        type: "bytes32",
-      },
-      {
-        internalType: "uint16",
-        name: "netuid",
-        type: "uint16",
-      },
+        name: "delegate",
+        type: "bytes32"
+      }
     ],
-    name: "addStake",
+    name: "addProxy",
     outputs: [],
-    stateMutability: "payable",
-    type: "function",
+    stateMutability: "nonpayable",
+    type: "function"
   },
   {
     inputs: [
       {
-        internalType: "bytes32",
-        name: "hotkey",
-        type: "bytes32",
+        "internalType": "bytes32",
+        "name": "hotkey",
+        "type": "bytes32"
       },
       {
-        internalType: "bytes32",
-        name: "coldkey",
-        type: "bytes32",
-      },
+        "internalType": "uint16",
+        "name": "netuid",
+        "type": "uint16"
+      }
     ],
-    name: "getStake",
-    outputs: [
+    name: "addStake",
+    outputs: [],
+    stateMutability: "payable",
+    type: "function"
+  },
+  {
+    inputs: [
       {
-        internalType: "uint64",
-        name: "",
-        type: "uint64",
-      },
+        "internalType": "bytes32",
+        "name": "delegate",
+        "type": "bytes32"
+      }
     ],
-    stateMutability: "view",
-    type: "function",
+    name: "removeProxy",
+    outputs: [],
+    stateMutability: "nonpayable",
+    type: "function"
   },
   {
     inputs: [
       {
         internalType: "bytes32",
         name: "hotkey",
-        type: "bytes32",
+        type: "bytes32"
       },
       {
         internalType: "uint256",
         name: "amount",
-        type: "uint256",
+        type: "uint256"
       },
       {
         internalType: "uint16",
         name: "netuid",
-        type: "uint16",
-      },
+        type: "uint16"
+      }
     ],
     name: "removeStake",
     outputs: [],
     stateMutability: "nonpayable",
-    type: "function",
-  },
+    type: "function"
+  }
 ];
 
 let address = "0x0000000000000000000000000000000000000801";
@@ -259,4 +259,38 @@ describe("Staking precompile", () => {
       expect(stake < before_stake);
     });
   });
+
+  it("Can add/remove proxy", async () => {
+    await usingEthApi(async (provider) => {
+      // add/remove are done in a single test case, because we can't use the same private/public key
+      // between substrate and EVM, but to test the remove part, we must predefine the proxy first.
+      // it makes `remove` being dependent on `add`, because we should use `addProxy` from contract 
+      // to prepare the proxy for `removeProxy` testing - the proxy is specified for the 
+      // caller/origin.
+
+      // first, check we don't have proxies
+      const publicKey = convertH160ToPublicKey(fundedEthWallet.address);
+      let proxies = await api.query.proxy.proxies(publicKey);
+      expect(proxies[0].length).to.be.eq(0);
+
+      // intialize the contract
+      const signer = new ethers.Wallet(fundedEthWallet.privateKey, provider);
+      const contract = new ethers.Contract(address, abi, signer);
+
+      // test "add"
+      let tx = await contract.addProxy(tk.bob.publicKey);
+      await tx.wait();
+
+      const [[{ delegate }]] = await api.query.proxy.proxies(publicKey);
+
+      expect(delegate.toHuman()).to.be.eq(tk.bob.address);
+
+      // test "remove" 
+      tx = await contract.removeProxy(tk.bob.publicKey);
+      await tx.wait();
+
+      proxies = await api.query.proxy.proxies(publicKey);
+      expect(proxies[0].length).to.be.eq(0);
+    });
+  });
 });
diff --git a/yarn.lock b/yarn.lock
index d121ae1..4c80c3d 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -638,6 +638,11 @@ chai-as-promised@^8.0.0:
   dependencies:
     check-error "^2.0.0"
 
+chai-bignumber@^3.1.0:
+  version "3.1.0"
+  resolved "https://registry.yarnpkg.com/chai-bignumber/-/chai-bignumber-3.1.0.tgz#e196456c760df21f0e124f6df922289ea15a7e4c"
+  integrity sha512-omxEc80jAU+pZwRmoWr3aEzeLad4JW3iBhLRQlgISvghBdIxrMT7mVAGsDz4WSyCkKowENshH2j9OABAhld7QQ==
+
 chai@^5.1.1:
   version "5.1.1"
   resolved "https://registry.yarnpkg.com/chai/-/chai-5.1.1.tgz#f035d9792a22b481ead1c65908d14bb62ec1c82c"