From 210b1b6c332d90b0a8b8004d52b8f1f13c2653f6 Mon Sep 17 00:00:00 2001
From: Aleksandr Ivlev <a.ivlev.om@gmail.com>
Date: Fri, 24 May 2024 16:28:41 +0300
Subject: [PATCH] Added balance check and consistancy check

---
 src/Lottery.test.ts | 50 +++++++++++++++++++++++++++++++++++++++++----
 src/constants.ts    |  2 +-
 2 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/src/Lottery.test.ts b/src/Lottery.test.ts
index 31b9a2b..42323a1 100644
--- a/src/Lottery.test.ts
+++ b/src/Lottery.test.ts
@@ -10,6 +10,7 @@ import {
 import { Lottery } from './Lottery';
 import { Ticket } from './Ticket';
 import { getEmpty2dMerkleMap } from './util';
+import { TICKET_PRICE } from './constants';
 
 /*
  * This file specifies how to test the `Add` example smart contract. It is safe to delete this file and replace
@@ -62,6 +63,28 @@ class StateManager {
     return [roundWitness, ticketRoundWitness];
   }
 
+  addTicket(
+    ticket: Ticket,
+    round: number
+  ): [MerkleMapWitness, MerkleMapWitness, MerkleMapWitness, Field] {
+    const [roundWitness, ticketRoundWitness] = this.getNextTicketWitenss(round);
+    const [bankWitness, bankValue] = this.getBankWitness(round);
+
+    this.roundTicketMap[round].set(
+      Field.from(this.lastTicketInRound[round]),
+      ticket.hash()
+    );
+    this.lastTicketInRound[round]++;
+    this.ticketMap.set(Field.from(round), this.roundTicketMap[round].getRoot());
+
+    this.bankMap.set(
+      Field.from(round),
+      bankValue.add(TICKET_PRICE.mul(ticket.amount).value)
+    );
+
+    return [roundWitness, ticketRoundWitness, bankWitness, bankValue];
+  }
+
   // Returns witness and value
   getBankWitness(round: number): [MerkleMapWitness, Field] {
     const bankWitness = this.bankMap.getWitness(Field.from(round));
@@ -81,7 +104,8 @@ describe('Add', () => {
     zkAppAddress: PublicKey,
     zkAppPrivateKey: PrivateKey,
     lottery: Lottery,
-    state: StateManager;
+    state: StateManager,
+    checkConsistancy: () => void;
 
   beforeAll(async () => {
     if (proofsEnabled) await Lottery.compile();
@@ -98,6 +122,17 @@ describe('Add', () => {
     zkAppAddress = zkAppPrivateKey.toPublicKey();
     lottery = new Lottery(zkAppAddress);
     state = new StateManager();
+
+    checkConsistancy = () => {
+      expect(lottery.ticketRoot.get()).toEqual(state.ticketMap.getRoot());
+      expect(lottery.ticketNullifier.get()).toEqual(
+        state.ticketNullifierMap.getRoot()
+      );
+      expect(lottery.bankRoot.get()).toEqual(state.bankMap.getRoot());
+      expect(lottery.roundResultRoot.get()).toEqual(
+        state.roundResultMap.getRoot()
+      );
+    };
   });
 
   async function localDeploy() {
@@ -115,11 +150,12 @@ describe('Add', () => {
 
     let curRound = 0;
 
+    const balanceBefore = Mina.getBalance(senderAccount);
+
     // Buy ticket
     const ticket = Ticket.random(senderAccount);
-    let [roundWitness, roundTicketWitness] =
-      state.getNextTicketWitenss(curRound);
-    let [bankWitness, bankValue] = state.getBankWitness(curRound);
+    let [roundWitness, roundTicketWitness, bankWitness, bankValue] =
+      state.addTicket(ticket, curRound);
     let tx = await Mina.transaction(senderAccount, async () => {
       await lottery.buyTicket(
         ticket,
@@ -133,6 +169,12 @@ describe('Add', () => {
     await tx.prove();
     await tx.sign([senderKey]).send();
 
+    const balanceAfter = Mina.getBalance(senderAccount);
+
+    expect(balanceBefore.sub(balanceAfter)).toEqual(TICKET_PRICE);
+
+    checkConsistancy();
+
     // Wait next round
 
     // Produce result
diff --git a/src/constants.ts b/src/constants.ts
index 361607e..817a8d4 100644
--- a/src/constants.ts
+++ b/src/constants.ts
@@ -2,7 +2,7 @@ import { UInt64 } from 'o1js';
 
 export const NUMBERS_IN_TICKET = 6;
 
-export const TICKET_PRICE = UInt64.from(10); // #TODO change to field in smartcontract
+export const TICKET_PRICE = UInt64.from(10 * 10 ** 9); // #TODO change to field in smartcontract
 export const BLOCK_PER_ROUND = 480; // Aproximate blocks per day
 
 export const SCORE_COEFFICIENTS = [1, 10, 100, 1000, 10000, 100000]; // Should be updated with apropriate probaility