fix: added test coverage to reproduce the drawing bug, contract fix

Author: jaybuidl

contracts/src/arbitration/SortitionModule.sol

@@ -289,11 +289,15 @@ contract SortitionModule is ISortitionModule {
         require(phase == Phase.drawing, "Wrong phase.");
         SortitionSumTree storage tree = sortitionSumTrees[_key];
 
-        uint256 treeIndex = 0;
+        if (tree.nodes[0] == 0) {
+            return address(0); // No jurors staked.
+        }
+
         uint256 currentDrawnNumber = uint256(keccak256(abi.encodePacked(randomNumber, _coreDisputeID, _voteID))) %
             tree.nodes[0];
 
         // While it still has children
+        uint256 treeIndex = 0;
         while ((tree.K * treeIndex) + 1 < tree.nodes.length) {
             for (uint256 i = 1; i <= tree.K; i++) {
                 // Loop over children.

contracts/src/arbitration/dispute-kits/DisputeKitClassic.sol

@@ -557,18 +557,8 @@ contract DisputeKitClassic is BaseDisputeKit, IEvidence {
     // *            Internal               * //
     // ************************************* //
 
-    /// @dev Checks that the chosen address satisfies certain conditions for being drawn.
-    /// @param _coreDisputeID ID of the dispute in the core contract.
-    /// @param _juror Chosen address.
-    /// @return Whether the address can be drawn or not.
-    function _postDrawCheck(uint256 _coreDisputeID, address _juror) internal view override returns (bool) {
-        (uint96 courtID, , , , ) = core.disputes(_coreDisputeID);
-        (, uint256 lockedAmountPerJuror, , , , , , , , ) = core.getRoundInfo(
-            _coreDisputeID,
-            core.getNumberOfRounds(_coreDisputeID) - 1
-        );
-        (uint256 staked, uint256 locked, ) = core.getJurorBalance(_juror, courtID);
-        (, , uint256 minStake, , , , ) = core.courts(courtID);
-        return staked >= locked + lockedAmountPerJuror && staked >= minStake;
+    /// @inheritdoc BaseDisputeKit
+    function _postDrawCheck(uint256 /*_coreDisputeID*/, address /*_juror*/) internal pure override returns (bool) {
+        return true;
     }
 }

contracts/test/arbitration/draw.ts

@@ -1,6 +1,6 @@
 import { anyValue } from "@nomicfoundation/hardhat-chai-matchers/withArgs";
 import { deployments, ethers, getNamedAccounts, network } from "hardhat";
-import { BigNumber } from "ethers";
+import { BigNumber, ContractTransaction, Wallet } from "ethers";
 import {
   PNK,
   KlerosCore,
@@ -52,11 +52,7 @@ describe("Draw Benchmark", async () => {
   let randomizer;
 
   beforeEach("Setup", async () => {
-    deployer = (await getNamedAccounts()).deployer;
-    relayer = (await getNamedAccounts()).relayer;
-
-    console.log("deployer:%s", deployer);
-    console.log("named accounts: %O", await getNamedAccounts());
+    ({ deployer, relayer } = await getNamedAccounts());
 
     await deployments.fixture(["Arbitration", "VeaMock"], {
       fallbackToGlobal: true,
@@ -70,9 +66,33 @@ describe("Draw Benchmark", async () => {
     rng = (await ethers.getContract("RandomizerRNG")) as RandomizerRNG;
     randomizer = (await ethers.getContract("RandomizerMock")) as RandomizerMock;
     sortitionModule = (await ethers.getContract("SortitionModule")) as SortitionModule;
+
+    // CourtId 2
+    const minStake = BigNumber.from(10).pow(20).mul(3); // 300 PNK
+    const alpha = 10000;
+    const feeForJuror = BigNumber.from(10).pow(17);
+    await core.createCourt(
+      1,
+      false,
+      minStake,
+      alpha,
+      feeForJuror,
+      256,
+      [0, 0, 0, 10], // evidencePeriod, commitPeriod, votePeriod, appealPeriod
+      ethers.utils.hexlify(5), // Extra data for sortition module will return the default value of K)
+      [1]
+    );
   });
 
-  it("Draw Benchmark", async () => {
+  interface SetStake {
+    (wallet: Wallet): Promise<void>;
+  }
+
+  interface ExpectFromDraw {
+    (drawTx: Promise<ContractTransaction>): Promise<void>;
+  }
+
+  const draw = async (setStake: SetStake, createDisputeCourtId: string, expectFromDraw: ExpectFromDraw) => {
     const arbitrationCost = ONE_TENTH_ETH.mul(3);
     const [bridger] = await ethers.getSigners();
 
@@ -90,7 +110,8 @@ describe("Draw Benchmark", async () => {
       expect(await pnk.balanceOf(wallet.address)).to.equal(ONE_THOUSAND_PNK.mul(10));
 
       await pnk.connect(wallet).approve(core.address, ONE_THOUSAND_PNK.mul(10), { gasLimit: 300000 });
-      await core.connect(wallet).setStake(1, ONE_THOUSAND_PNK.mul(10), { gasLimit: 5000000 });
+
+      await setStake(wallet);
     }
 
     // Create a dispute
@@ -114,7 +135,7 @@ describe("Draw Benchmark", async () => {
           templateId: 0,
           templateUri: "",
           choices: 2,
-          extraData: "0x00",
+          extraData: `0x000000000000000000000000000000000000000000000000000000000000000${createDisputeCourtId}0000000000000000000000000000000000000000000000000000000000000003`,
         },
         { value: arbitrationCost }
       );
@@ -131,12 +152,72 @@ describe("Draw Benchmark", async () => {
     await randomizer.relay(rng.address, 0, ethers.utils.randomBytes(32));
     await sortitionModule.passPhase(); // Generating -> Drawing
 
-    await expect(core.draw(0, 1000, { gasLimit: 1000000 }))
-      .to.emit(core, "Draw")
-      .withArgs(anyValue, 0, 0, 0)
-      .to.emit(core, "Draw")
-      .withArgs(anyValue, 0, 0, 1)
-      .to.emit(core, "Draw")
-      .withArgs(anyValue, 0, 0, 2);
+    await expectFromDraw(core.draw(0, 1000, { gasLimit: 1000000 }));
+  };
+
+  it("Stakes in parent court and should draw jurors in parent court", async () => {
+    const setStake = async (wallet: Wallet) => {
+      await core.connect(wallet).setStake(1, ONE_THOUSAND_PNK.mul(5), { gasLimit: 5000000 });
+    };
+
+    const expectFromDraw = async (drawTx: Promise<ContractTransaction>) => {
+      await expect(drawTx)
+        .to.emit(core, "Draw")
+        .withArgs(anyValue, 0, 0, 0)
+        .to.emit(core, "Draw")
+        .withArgs(anyValue, 0, 0, 1)
+        .to.emit(core, "Draw")
+        .withArgs(anyValue, 0, 0, 2);
+    };
+
+    await draw(setStake, "1", expectFromDraw);
+  });
+
+  it("Stakes in parent court and should draw nobody in subcourt", async () => {
+    const setStake = async (wallet: Wallet) => {
+      await core.connect(wallet).setStake(1, ONE_THOUSAND_PNK.mul(5), { gasLimit: 5000000 });
+    };
+
+    const expectFromDraw = async (drawTx: Promise<ContractTransaction>) => {
+      await expect(drawTx).to.not.emit(core, "Draw");
+    };
+
+    await draw(setStake, "2", expectFromDraw);
+  });
+
+  it("Stakes in subcourt and should draw jurors in parent court", async () => {
+    const setStake = async (wallet: Wallet) => {
+      await core.connect(wallet).setStake(2, ONE_THOUSAND_PNK.mul(5), { gasLimit: 5000000 });
+    };
+
+    const expectFromDraw = async (drawTx: Promise<ContractTransaction>) => {
+      await expect(drawTx)
+        .to.emit(core, "Draw")
+        .withArgs(anyValue, 0, 0, 0)
+        .to.emit(core, "Draw")
+        .withArgs(anyValue, 0, 0, 1)
+        .to.emit(core, "Draw")
+        .withArgs(anyValue, 0, 0, 2);
+    };
+
+    await draw(setStake, "1", expectFromDraw);
+  });
+
+  it("Stakes in subcourt and should draw jurors in subcourt", async () => {
+    const setStake = async (wallet: Wallet) => {
+      await core.connect(wallet).setStake(2, ONE_THOUSAND_PNK.mul(5), { gasLimit: 5000000 });
+    };
+
+    const expectFromDraw = async (drawTx: Promise<ContractTransaction>) => {
+      await expect(drawTx)
+        .to.emit(core, "Draw")
+        .withArgs(anyValue, 0, 0, 0)
+        .to.emit(core, "Draw")
+        .withArgs(anyValue, 0, 0, 1)
+        .to.emit(core, "Draw")
+        .withArgs(anyValue, 0, 0, 2);
+    };
+
+    await draw(setStake, "2", expectFromDraw);
   });
 });

contracts/test/arbitration/unstake.ts

@@ -31,10 +31,6 @@ describe("Unstake juror", async () => {
 
   beforeEach("Setup", async () => {
     ({ deployer } = await getNamedAccounts());
-
-    console.log("deployer:%s", deployer);
-    console.log("named accounts: %O", await getNamedAccounts());
-
     await deployments.fixture(["Arbitration"], {
       fallbackToGlobal: true,
       keepExistingDeployments: false,

contracts/test/integration/index.ts

@@ -43,10 +43,6 @@ describe("Integration tests", async () => {
 
   beforeEach("Setup", async () => {
     ({ deployer } = await getNamedAccounts());
-
-    console.log("deployer:%s", deployer);
-    console.log("named accounts: %O", await getNamedAccounts());
-
     await deployments.fixture(["Arbitration", "VeaMock"], {
       fallbackToGlobal: true,
       keepExistingDeployments: false,