From 576ba47938380366a622d0a90d463b3f21f7ce20 Mon Sep 17 00:00:00 2001
From: poma <semenov.roma@gmail.com>
Date: Sat, 21 Aug 2021 15:46:17 +0300
Subject: [PATCH] use waffle

---
 src/utils.js      |  9 --------
 test/full.test.js | 53 ++++++++++++++++++++++++-----------------------
 2 files changed, 27 insertions(+), 35 deletions(-)

diff --git a/src/utils.js b/src/utils.js
index 87b9e5f..33ab330 100644
--- a/src/utils.js
+++ b/src/utils.js
@@ -75,13 +75,6 @@ function shuffle(array) {
 
   return array
 }
-async function takeSnapshot() {
-  return await ethers.provider.send('evm_snapshot', [])
-}
-
-async function revertSnapshot(id) {
-  await ethers.provider.send('evm_revert', [id])
-}
 
 module.exports = {
   FIELD_SIZE,
@@ -91,7 +84,5 @@ module.exports = {
   poseidonHash,
   poseidonHash2,
   getExtDataHash,
-  takeSnapshot,
-  revertSnapshot,
   shuffle,
 }
diff --git a/test/full.test.js b/test/full.test.js
index 61d209b..4790fb3 100644
--- a/test/full.test.js
+++ b/test/full.test.js
@@ -1,9 +1,9 @@
 const hre = require('hardhat')
-const ethers = hre.ethers
-const { expect, should } = require('chai')
-should()
+const { ethers, waffle } = hre
+const { loadFixture } = waffle
+const { expect } = require('chai')
 
-const { poseidonHash2, toFixedHex, takeSnapshot, revertSnapshot } = require('../src/utils')
+const { poseidonHash2, toFixedHex } = require('../src/utils')
 const Utxo = require('../src/utxo')
 
 const MERKLE_TREE_HEIGHT = 5
@@ -14,30 +14,29 @@ const { Keypair } = require('../src/keypair')
 
 describe('TornadoPool', function () {
   this.timeout(20000)
-  let snapshotId, sender
-  /** @type {TornadoPool} */
-  let tornadoPool
 
-  /* prettier-ignore */
-  before(async function () {
-    ;[sender] = await ethers.getSigners()
+  async function deploy(contractName, ...args) {
+    const Factory = await ethers.getContractFactory(contractName)
+    const instance = await Factory.deploy(...args)
+    await instance.deployed()
 
-    const Verifier2 = await ethers.getContractFactory('Verifier2')
-    const verifier2 = await Verifier2.deploy()
-    await verifier2.deployed()
+    return instance
+  }
 
-    const Verifier16 = await ethers.getContractFactory('Verifier16')
-    const verifier16 = await Verifier16.deploy()
-    await verifier16.deployed()
+  async function fixture() {
+    const [deployer] = await ethers.getSigners()
+
+    const verifier2 = await deploy('Verifier2')
+
+    const verifier16 = await deploy('Verifier16')
 
     const tree = new MerkleTree(MERKLE_TREE_HEIGHT, [], { hashFunction: poseidonHash2 })
     const root = await tree.root()
+    /** @type {TornadoPool} */
+    const tornadoPool = await deploy('TornadoPool', verifier2.address, verifier16.address, toFixedHex(root))
 
-    const Pool = await ethers.getContractFactory('TornadoPool')
-    tornadoPool = await Pool.deploy(verifier2.address, verifier16.address, toFixedHex(root))
-
-    snapshotId = await takeSnapshot()
-  })
+    return { tornadoPool, deployer }
+  }
 
   it('encrypt -> decrypt should work', () => {
     const data = Buffer.from([0xff, 0xaa, 0x00, 0x01])
@@ -49,6 +48,7 @@ describe('TornadoPool', function () {
   })
 
   it('constants check', async () => {
+    const { tornadoPool } = await loadFixture(fixture)
     const maxFee = await tornadoPool.MAX_FEE()
     const maxExtAmount = await tornadoPool.MAX_EXT_AMOUNT()
     const fieldSize = await tornadoPool.FIELD_SIZE()
@@ -57,6 +57,9 @@ describe('TornadoPool', function () {
   })
 
   it('should register and deposit', async function () {
+    let { tornadoPool } = await loadFixture(fixture)
+    const sender = (await ethers.getSigners())[1]
+
     // Alice deposits into tornado pool
     const aliceDepositAmount = 1e7
     const aliceDepositUtxo = new Utxo({ amount: aliceDepositAmount })
@@ -115,6 +118,8 @@ describe('TornadoPool', function () {
   })
 
   it('should deposit, transact and withdraw', async function () {
+    const { tornadoPool } = await loadFixture(fixture)
+
     // Alice deposits into tornado pool
     const aliceDepositAmount = 1e7
     const aliceDepositUtxo = new Utxo({ amount: aliceDepositAmount })
@@ -162,11 +167,7 @@ describe('TornadoPool', function () {
   })
 
   it('should work with 16 inputs', async function () {
+    const { tornadoPool } = await loadFixture(fixture)
     await transaction({ tornadoPool, inputs: [new Utxo(), new Utxo(), new Utxo()] })
   })
-
-  afterEach(async () => {
-    await revertSnapshot(snapshotId)
-    snapshotId = await takeSnapshot()
-  })
 })