test(core): add test cases for processMessage()

privacy-scaling-explorations · Dec 12, 2023 · 3c0891c · 3c0891c
1 parent bd2a66d
commit 3c0891c
Show file tree

Hide file tree

Showing 2 changed files with 121 additions and 6 deletions.
diff --git a/core/ts/__tests__/MaciState.test.ts b/core/ts/__tests__/MaciState.test.ts
@@ -109,7 +109,7 @@ describe("MaciState", function () {
       accumulatorQueue.mergeSubRoots(0);
       accumulatorQueue.merge(treeDepths.messageTreeDepth);
 
-      expect(accumulatorQueue.getRoot(treeDepths.messageTreeDepth).toString()).to.eq(msgTree.root.toString());
+      expect(accumulatorQueue.getRoot(treeDepths.messageTreeDepth).toString()).to.eq(poll.messageTree.root.toString());
     });
 
     it("packProcessMessageSmallVals and unpackProcessMessageSmallVals", () => {

diff --git a/core/ts/__tests__/ProcessMessage.test.ts b/core/ts/__tests__/ProcessMessage.test.ts
@@ -1,5 +1,120 @@
-/* eslint-disable */
-describe("Message processing", () => {
+import { expect } from "chai";
+
+import { PCommand, Message, Keypair, PubKey } from "maci-domainobjs";
+import { Signature } from "maci-crypto";
+
+import { STATE_TREE_DEPTH, MaciState, Poll } from "../";
+
+const voiceCreditBalance = BigInt(100);
+
+const duration = 30;
+const maxValues = {
+  maxUsers: 25,
+  maxMessages: 25,
+  maxVoteOptions: 25,
+};
+
+const treeDepths = {
+  intStateTreeDepth: 2,
+  messageTreeDepth: 3,
+  messageTreeSubDepth: 2,
+  voteOptionTreeDepth: 4,
+};
+
+const messageBatchSize = 25;
+
+class TestHarness {
+  maciState = new MaciState(STATE_TREE_DEPTH);
+  coordinatorKeypair = new Keypair();
+  poll: Poll;
+  pollId: number;
+  users: Keypair[] = [];
+  stateIndices = new Map<Keypair, number>();
+
+  constructor() {
+    this.pollId = this.maciState.deployPoll(
+      duration,
+      BigInt(Math.floor(Date.now() / 1000) + duration),
+      maxValues,
+      treeDepths,
+      messageBatchSize,
+      this.coordinatorKeypair,
+    );
+    this.poll = this.maciState.polls[this.pollId];
+  }
+
+  createUsers = (numUsers: number): Keypair[] => {
+    for (let i = 0; i < numUsers; i++) {
+      const user = new Keypair();
+      this.users.push(user);
+      const stateIndex = this.signup(user);
+      this.stateIndices.set(user, stateIndex);
+    }
+    return this.users;
+  };
+
+  signup = (user: Keypair): number => {
+    const timestamp = BigInt(Math.floor(Date.now() / 1000));
+    const stateIndex = this.maciState.signUp(user.pubKey, voiceCreditBalance, timestamp);
+    return stateIndex;
+  };
+
+  vote = (user: Keypair, stateIndex: number, voteOptionIndex: bigint, voteWeight: bigint, nonce: bigint): void => {
+    const { command, signature } = this.createCommand(user, stateIndex, voteOptionIndex, voteWeight, nonce);
+
+    const { message, encPubKey } = this.createMessage(command, signature, this.coordinatorKeypair);
+
+    this.poll.publishMessage(message, encPubKey);
+  };
+
+  createMessage = (
+    command: PCommand,
+    signature: Signature,
+    coordinatorKeypair: Keypair,
+  ): { message: Message; encPubKey: PubKey } => {
+    const ecdhKeypair = new Keypair();
+    const sharedKey = Keypair.genEcdhSharedKey(ecdhKeypair.privKey, coordinatorKeypair.pubKey);
+    const message = command.encrypt(signature, sharedKey);
+    return { message, encPubKey: ecdhKeypair.pubKey };
+  };
+
+  createCommand = (
+    user: Keypair,
+    stateIndex: number,
+    voteOptionIndex: bigint,
+    voteWeight: bigint,
+    nonce: bigint,
+  ): { command: PCommand; signature: Signature } => {
+    const command = new PCommand(
+      BigInt(stateIndex),
+      user.pubKey,
+      voteOptionIndex,
+      voteWeight,
+      nonce,
+      BigInt(this.pollId),
+    );
+
+    const signature = command.sign(user.privKey);
+
+    return { command, signature };
+  };
+
+  finalizePoll = (): void => {
+    this.poll.processMessages(this.pollId);
+    this.poll.tallyVotes();
+  };
+
+  getStateIndex = (user: Keypair): number => {
+    return this.stateIndices.get(user);
+  };
+}
+
+describe("Message processing", function () {
+  this.timeout(30000);
+
+  let testHarness: TestHarness;
+  let poll: Poll;
+
   describe("Process a batch of messages", () => {
     beforeEach(async () => {
       testHarness = new TestHarness();
@@ -20,7 +135,7 @@ describe("Message processing", () => {
       expect(poll.tallyResult[0]).to.eq(BigInt(9));
     });
 
-    it("processMessage() should not process messages twice", async () => {
+    it.only("processMessage() should not process messages twice", async () => {
       const voteOptionIndex = BigInt(0);
       const voteWeight = BigInt(9);
       const nonce = BigInt(1);
@@ -123,7 +238,7 @@ describe("Message processing", () => {
       expect(tallyResult).to.eq(expectedTallyResult);
     });
 
-    it("processMessage() should not process a message with an incorrect signature", async () => {
+    it.only("processMessage() should not process a message with an incorrect signature", async () => {
       const voteOptionIndex = BigInt(0);
       const voteWeight = BigInt(9);
       const nonce = BigInt(1);
@@ -166,7 +281,7 @@ describe("Message processing", () => {
       expect(tallyResult).to.eq(expectedTallyResult);
     });
 
-    it("processMessage() should not process a message with an invalid coordinator key", async () => {
+    it.only("processMessage() should not process a message with an invalid coordinator key", async () => {
       const voteOptionIndex = BigInt(0);
       const voteWeight = BigInt(9);
       const nonce = BigInt(1);