From 04759deb4f74c92bb92292a266d1311a155af179 Mon Sep 17 00:00:00 2001 From: Daehyun Paik Date: Mon, 18 Dec 2023 01:05:30 +0100 Subject: [PATCH] refactor(core): export `TestHarness` from core test utils --- core/ts/__tests__/ProcessMessages.test.ts | 109 +-------------- core/ts/__tests__/utils/MaciTestHarness.ts | 155 +++++++++++++++++++++ core/ts/index.ts | 2 + 3 files changed, 161 insertions(+), 105 deletions(-) create mode 100644 core/ts/__tests__/utils/MaciTestHarness.ts diff --git a/core/ts/__tests__/ProcessMessages.test.ts b/core/ts/__tests__/ProcessMessages.test.ts index d43e7e1a3a..9f8799c0f9 100644 --- a/core/ts/__tests__/ProcessMessages.test.ts +++ b/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(); - - 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 diff --git a/core/ts/__tests__/utils/MaciTestHarness.ts b/core/ts/__tests__/utils/MaciTestHarness.ts new file mode 100644 index 0000000000..8351191f24 --- /dev/null +++ b/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(); + + /** + * 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); + }; +} diff --git a/core/ts/index.ts b/core/ts/index.ts index 9e5bfa1e0b..4d3ca3bab3 100644 --- a/core/ts/index.ts +++ b/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";