test: updated with token whitelist coverage

Author: jaybuidl

contracts/src/test/DisputeKitGatedMock.sol

@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.24;
+
+import "../arbitration/dispute-kits/DisputeKitGated.sol";
+
+/// @title KlerosCoreMock
+/// KlerosCore with view functions to use in Foundry tests.
+contract DisputeKitGatedMock is DisputeKitGated {
+    function extraDataToTokenInfo(
+        bytes memory _extraData
+    ) public pure returns (address tokenGate, bool isERC1155, uint256 tokenId) {
+        (tokenGate, isERC1155, tokenId) = _extraDataToTokenInfo(_extraData);
+    }
+}

contracts/test/arbitration/dispute-kit-gated.ts

@@ -19,7 +19,18 @@ import { deployUpgradable } from "../../deploy/utils/deployUpgradable";
 /* eslint-disable no-unused-vars */
 /* eslint-disable no-unused-expressions */ // https://github.com/standard/standard/issues/690#issuecomment-278533482
 
+/**
+ * Test suite for DisputeKitGated - a dispute kit that requires jurors to hold
+ * specific tokens (ERC20, ERC721, or ERC1155) to participate in disputes.
+ *
+ * Tests cover:
+ * - Token whitelist management and access control
+ * - Error handling for unsupported tokens
+ * - Token gating functionality for different token types
+ * - Integration between whitelist and dispute creation
+ */
 describe("DisputeKitGated", async () => {
+  // Constants for token amounts
   const ONE_THOUSAND_PNK = 10n ** 21n;
   const thousandPNK = (amount: BigNumberish) => toBigInt(amount) * ONE_THOUSAND_PNK;
   const PNK = (amount: BigNumberish) => toBigInt(amount) * 10n ** 18n;
@@ -35,8 +46,8 @@ describe("DisputeKitGated", async () => {
   let rng: IncrementalNG;
   let nft721: TestERC721;
   let nft1155: TestERC1155;
+  let gatedDisputeKitID: number;
   const RANDOM = 424242n;
-  const GATED_DK_ID = 3;
   const TOKEN_ID = 888;
   const minStake = PNK(200);
 
@@ -62,8 +73,8 @@ describe("DisputeKitGated", async () => {
       log: true,
     });
     await core.addNewDisputeKit(deploymentResult.address);
-    const disputeKitGatedID = (await core.getDisputeKitsLength()) - 1n;
-    await core.enableDisputeKits(Courts.GENERAL, [disputeKitGatedID], true); // enable disputeKitGated on the General Court
+    gatedDisputeKitID = Number((await core.getDisputeKitsLength()) - 1n);
+    await core.enableDisputeKits(Courts.GENERAL, [gatedDisputeKitID], true); // enable disputeKitGated on the General Court
 
     disputeKitGated = await ethers.getContract<DisputeKitGatedMock>("DisputeKitGatedMock");
 
@@ -84,8 +95,14 @@ describe("DisputeKitGated", async () => {
     await deployERC1155(hre, deployer, "TestERC1155", "Nft1155");
     nft1155 = await ethers.getContract<TestERC1155>("Nft1155");
     await nft1155.mint(deployer, TOKEN_ID, 1, "0x00");
+
+    // Whitelist all tokens by default to make existing tests work
+    await whitelistTokens([dai.target, nft721.target, nft1155.target], true);
   });
 
+  // Helper Functions
+
+  /** Encodes extra data for dispute creation with token gating parameters */
   const encodeExtraData = (
     courtId: number,
     minJurors: BigNumberish,
@@ -103,6 +120,29 @@ describe("DisputeKitGated", async () => {
     );
   };
 
+  /** Adds or removes tokens from the whitelist */
+  const whitelistTokens = async (tokens: (string | Addressable)[], supported: boolean = true) => {
+    const tokenAddresses = tokens.map((token) => (typeof token === "string" ? token : token.toString()));
+    return disputeKitGated.changeSupportedTokens(tokenAddresses, supported);
+  };
+
+  /** Creates a dispute with the specified token gate */
+  const createDisputeWithToken = async (
+    token: string | Addressable,
+    isERC1155: boolean = false,
+    tokenId: BigNumberish = 0
+  ) => {
+    const extraData = encodeExtraData(Courts.GENERAL, 3, gatedDisputeKitID, token, isERC1155, tokenId);
+    const arbitrationCost = await core["arbitrationCost(bytes)"](extraData);
+    return core["createDispute(uint256,bytes)"](2, extraData, { value: arbitrationCost });
+  };
+
+  /** Asserts whether a token is supported or not */
+  const expectTokenSupported = async (token: string | Addressable, supported: boolean) => {
+    const tokenAddress = typeof token === "string" ? token : token.toString();
+    expect(await disputeKitGated.supportedTokens(tokenAddress)).to.equal(supported);
+  };
+
   const stakeAndDraw = async (
     courtId: number,
     minJurors: BigNumberish,
@@ -159,10 +199,122 @@ describe("DisputeKitGated", async () => {
     return core.draw(disputeId, 70, { gasLimit: 10000000 });
   };
 
+  describe("Token Whitelist Management", async () => {
+    describe("changeSupportedTokens function", async () => {
+      it("Should allow owner to whitelist single token", async () => {
+        await whitelistTokens([dai.target], true);
+        await expectTokenSupported(dai.target, true);
+      });
+
+      it("Should allow owner to whitelist multiple tokens", async () => {
+        await whitelistTokens([dai.target, nft721.target, nft1155.target], true);
+        await expectTokenSupported(dai.target, true);
+        await expectTokenSupported(nft721.target, true);
+        await expectTokenSupported(nft1155.target, true);
+      });
+
+      it("Should allow owner to remove single token from whitelist", async () => {
+        await whitelistTokens([dai.target], true);
+        await expectTokenSupported(dai.target, true);
+
+        await whitelistTokens([dai.target], false);
+        await expectTokenSupported(dai.target, false);
+      });
+
+      it("Should allow owner to remove multiple tokens from whitelist", async () => {
+        await whitelistTokens([dai.target, nft721.target], true);
+        await expectTokenSupported(dai.target, true);
+        await expectTokenSupported(nft721.target, true);
+
+        await whitelistTokens([dai.target, nft721.target], false);
+        await expectTokenSupported(dai.target, false);
+        await expectTokenSupported(nft721.target, false);
+      });
+
+      it("Should handle mixed operations (add some, remove some)", async () => {
+        await whitelistTokens([dai.target, nft721.target], true);
+        await expectTokenSupported(dai.target, true);
+        await expectTokenSupported(nft721.target, true);
+
+        await whitelistTokens([dai.target], false);
+        await whitelistTokens([nft1155.target], true);
+
+        await expectTokenSupported(dai.target, false);
+        await expectTokenSupported(nft721.target, true);
+        await expectTokenSupported(nft1155.target, true);
+      });
+
+      it("Should handle duplicate operations correctly", async () => {
+        // Whitelist token twice - should not revert
+        await whitelistTokens([dai.target], true);
+        await whitelistTokens([dai.target], true);
+        await expectTokenSupported(dai.target, true);
+
+        // Remove token twice - should not revert
+        await whitelistTokens([dai.target], false);
+        await whitelistTokens([dai.target], false);
+        await expectTokenSupported(dai.target, false);
+      });
+    });
+
+    describe("Access Control", async () => {
+      it("Should revert when non-owner tries to change supported tokens", async () => {
+        await expect(disputeKitGated.connect(juror1).changeSupportedTokens([dai.target], true)).to.be.reverted;
+      });
+
+      it("Should revert when non-owner tries to remove supported tokens", async () => {
+        // First whitelist as owner
+        await whitelistTokens([dai.target], true);
+
+        // Then try to remove as non-owner
+        await expect(disputeKitGated.connect(juror1).changeSupportedTokens([dai.target], false)).to.be.reverted;
+      });
+    });
+  });
+
+  describe("Error Handling - Unsupported Tokens", async () => {
+    it("Should revert with TokenNotSupported when creating dispute with unsupported ERC20", async () => {
+      await whitelistTokens([dai.target], false);
+
+      await expect(createDisputeWithToken(dai.target))
+        .to.be.revertedWithCustomError(disputeKitGated, "TokenNotSupported")
+        .withArgs(dai.target);
+    });
+
+    it("Should revert with TokenNotSupported when creating dispute with unsupported ERC721", async () => {
+      await whitelistTokens([nft721.target], false);
+
+      await expect(createDisputeWithToken(nft721.target))
+        .to.be.revertedWithCustomError(disputeKitGated, "TokenNotSupported")
+        .withArgs(nft721.target);
+    });
+
+    it("Should revert with TokenNotSupported when creating dispute with unsupported ERC1155", async () => {
+      await whitelistTokens([nft1155.target], false);
+
+      await expect(createDisputeWithToken(nft1155.target, true, TOKEN_ID))
+        .to.be.revertedWithCustomError(disputeKitGated, "TokenNotSupported")
+        .withArgs(nft1155.target);
+    });
+
+    it("Should allow dispute creation after token is whitelisted", async () => {
+      await whitelistTokens([dai.target], false);
+
+      await expect(createDisputeWithToken(dai.target)).to.be.revertedWithCustomError(
+        disputeKitGated,
+        "TokenNotSupported"
+      );
+
+      await whitelistTokens([dai.target], true);
+
+      await expect(createDisputeWithToken(dai.target)).to.not.be.reverted;
+    });
+  });
+
   describe("When gating with DAI token", async () => {
     it("Should draw no juror if they don't have any DAI balance", async () => {
       const nbOfJurors = 15n;
-      const tx = await stakeAndDraw(Courts.GENERAL, nbOfJurors, GATED_DK_ID, dai.target, false, 0).then((tx) =>
+      const tx = await stakeAndDraw(Courts.GENERAL, nbOfJurors, gatedDisputeKitID, dai.target, false, 0).then((tx) =>
         tx.wait()
       );
 
@@ -176,7 +328,7 @@ describe("DisputeKitGated", async () => {
       dai.transfer(juror1.address, 1);
 
       const nbOfJurors = 15n;
-      const tx = await stakeAndDraw(Courts.GENERAL, nbOfJurors, GATED_DK_ID, dai.target, false, 0).then((tx) =>
+      const tx = await stakeAndDraw(Courts.GENERAL, nbOfJurors, gatedDisputeKitID, dai.target, false, 0).then((tx) =>
         tx.wait()
       );
 
@@ -209,7 +361,7 @@ describe("DisputeKitGated", async () => {
   describe("When gating with ERC721 token", async () => {
     it("Should draw no juror if they don't own the ERC721 token", async () => {
       const nbOfJurors = 15n;
-      const tx = await stakeAndDraw(Courts.GENERAL, nbOfJurors, GATED_DK_ID, nft721.target, false, 0).then((tx) =>
+      const tx = await stakeAndDraw(Courts.GENERAL, nbOfJurors, gatedDisputeKitID, nft721.target, false, 0).then((tx) =>
         tx.wait()
       );
 
@@ -223,7 +375,7 @@ describe("DisputeKitGated", async () => {
       await nft721.safeMint(juror2.address);
 
       const nbOfJurors = 15n;
-      const tx = await stakeAndDraw(Courts.GENERAL, nbOfJurors, GATED_DK_ID, nft721.target, false, 0).then((tx) =>
+      const tx = await stakeAndDraw(Courts.GENERAL, nbOfJurors, gatedDisputeKitID, nft721.target, false, 0).then((tx) =>
         tx.wait()
       );
 
@@ -256,7 +408,7 @@ describe("DisputeKitGated", async () => {
   describe("When gating with ERC1155 token", async () => {
     it("Should draw no juror if they don't own the ERC1155 token", async () => {
       const nbOfJurors = 15n;
-      const tx = await stakeAndDraw(Courts.GENERAL, nbOfJurors, GATED_DK_ID, nft1155.target, true, TOKEN_ID).then(
+      const tx = await stakeAndDraw(Courts.GENERAL, nbOfJurors, gatedDisputeKitID, nft1155.target, true, TOKEN_ID).then(
         (tx) => tx.wait()
       );
 
@@ -270,7 +422,7 @@ describe("DisputeKitGated", async () => {
       await nft1155.mint(juror2.address, TOKEN_ID, 1, "0x00");
 
       const nbOfJurors = 15n;
-      const tx = await stakeAndDraw(Courts.GENERAL, nbOfJurors, GATED_DK_ID, nft1155.target, true, TOKEN_ID).then(
+      const tx = await stakeAndDraw(Courts.GENERAL, nbOfJurors, gatedDisputeKitID, nft1155.target, true, TOKEN_ID).then(
         (tx) => tx.wait()
       );
 
@@ -299,4 +451,61 @@ describe("DisputeKitGated", async () => {
       ]);
     });
   });
+
+  describe("Whitelist Integration Tests", async () => {
+    it("Should allow new disputes after whitelisting a token", async () => {
+      // Whitelist DAI token
+      await whitelistTokens([dai.target], true);
+
+      // Transfer DAI to juror1 for token gating
+      await dai.transfer(juror1.address, 1);
+
+      const nbOfJurors = 3n;
+      const tx = await stakeAndDraw(Courts.GENERAL, nbOfJurors, gatedDisputeKitID, dai.target, false, 0).then((tx) =>
+        tx.wait()
+      );
+
+      // Verify dispute was created and juror drawn
+      const drawLogs =
+        tx?.logs.filter((log: any) => log.fragment?.name === "Draw" && log.address === core.target) || [];
+      expect(drawLogs).to.have.length(nbOfJurors);
+    });
+
+    it("Should prevent new disputes after removing token from whitelist", async () => {
+      await whitelistTokens([dai.target], true);
+      await dai.transfer(juror1.address, 1);
+
+      // Create first dispute (should work)
+      await expect(createDisputeWithToken(dai.target)).to.not.be.reverted;
+
+      // Remove token from whitelist
+      await whitelistTokens([dai.target], false);
+
+      // Try to create another dispute (should fail)
+      await expect(createDisputeWithToken(dai.target))
+        .to.be.revertedWithCustomError(disputeKitGated, "TokenNotSupported")
+        .withArgs(dai.target);
+    });
+
+    it("Should maintain whitelist state correctly across multiple operations", async () => {
+      const tokens = [dai.target, nft721.target, nft1155.target];
+
+      // All tokens should already be supported from the main setup
+      for (const token of tokens) {
+        await expectTokenSupported(token, true);
+      }
+
+      // Remove middle token
+      await whitelistTokens([nft721.target], false);
+      await expectTokenSupported(dai.target, true);
+      await expectTokenSupported(nft721.target, false);
+      await expectTokenSupported(nft1155.target, true);
+
+      // Re-add middle token
+      await whitelistTokens([nft721.target], true);
+      for (const token of tokens) {
+        await expectTokenSupported(token, true);
+      }
+    });
+  });
 });