refactor(core): export TestHarness from core test utils

core/ts/__tests__/ProcessMessages.test.ts

@@ -1,111 +1,10 @@
+import { AssertionError } from "assert";
 import { expect } from "chai";
 
-import { Signature } from "maci-crypto";
-import { PCommand, Message, Keypair, PubKey } from "maci-domainobjs";
-
-import { STATE_TREE_DEPTH, MaciState, Poll } from "..";
-import { AssertionError } from "assert";
+import { Keypair } from "maci-domainobjs";
 
-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);
-  };
-}
+import { Poll } from "..";
+import { TestHarness } from "./utils/MaciTestHarness";
 
 describe("Poll message processing and validation", function () {
   // set timeout to 30 seconds

core/ts/__tests__/utils/MaciTestHarness.ts

@@ -0,0 +1,155 @@
+import { Signature } from "maci-crypto";
+import { PCommand, Message, Keypair, PubKey } from "maci-domainobjs";
+
+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;
+
+/**
+ * A test harness for the MACI contract.
+ */
+export class TestHarness {
+  maciState = new MaciState(STATE_TREE_DEPTH);
+  coordinatorKeypair = new Keypair();
+  poll: Poll;
+  pollId: number;
+  users: Keypair[] = [];
+  stateIndices = new Map<Keypair, number>();
+
+  /**
+   * Constructs a new TestHarness object.
+   */
+  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];
+  }
+
+  /**
+   * Creates a number of users and signs them up to the MACI state tree.
+   * @param numUsers - The number of users to create.
+   * @returns The keypairs of the newly created users.
+   */
+  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;
+  };
+
+  /**
+   * Signs up a user to the MACI state tree.
+   * @param user - The keypair of the user.
+   * @returns The index of the newly signed-up user in the state tree.
+   */
+  signup = (user: Keypair): number => {
+    const timestamp = BigInt(Math.floor(Date.now() / 1000));
+    const stateIndex = this.maciState.signUp(user.pubKey, voiceCreditBalance, timestamp);
+    return stateIndex;
+  };
+
+  /**
+   * Publishes a message to the MACI poll instance.
+   * @param user - The keypair of the user.
+   * @param stateIndex - The index of the user in the state tree.
+   * @param voteOptionIndex - The index of the vote option.
+   * @param voteWeight - The weight of the vote.
+   * @param nonce - The nonce of the vote.
+   */
+  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);
+  };
+
+  /**
+   * Creates a message from a command and signature.
+   * @param command - The command to be encrypted.
+   * @param signature - The signature of the command signer.
+   * @param coordinatorKeypair - The keypair of the MACI round coordinator.
+   * @returns The message and the ephemeral public key used to encrypt the message.
+   */
+  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 };
+  };
+
+  /**
+   * Creates a command and signature.
+   * @param user - The keypair of the user.
+   * @param stateIndex - The index of the user in the state tree.
+   * @param voteOptionIndex - The index of the vote option.
+   * @param voteWeight - The weight of the vote.
+   * @param nonce - The nonce of the vote.
+   * @returns The command and signature of the command.
+   */
+  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 };
+  };
+
+  /**
+   * Finalizes the poll.
+   * This processes all messages and tallies the votes.
+   * This should be called after all votes have been cast.
+   */
+  finalizePoll = (): void => {
+    this.poll.processMessages(this.pollId);
+    this.poll.tallyVotes();
+  };
+
+  /**
+   * Returns the state index of a signed-up user.
+   * @param user - The keypair of the user.
+   * @returns The state index of the user.
+   */
+  getStateIndex = (user: Keypair): number => {
+    return this.stateIndices.get(user);
+  };
+}

core/ts/index.ts

@@ -19,3 +19,5 @@ export {
   MESSAGE_TREE_ARITY,
   VOTE_OPTION_TREE_ARITY,
 } from "./utils/constants";
+
+export { TestHarness } from "./__tests__/utils/MaciTestHarness";