fix snapshot.restore when used multiple times

OpenZeppelin · Oct 5, 2021 · 9da815d · 9da815d
1 parent 2fec6fe
commit 9da815d
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 8 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # Changelog
 
+## 0.5.15 (unreleased)
+
+ * Fix `snapshot.restore` when used multiple times.
+
 ## 0.5.14 (2021-10-04)
 
  * Add missing export of `snapshot` module.

diff --git a/src/snapshot.js b/src/snapshot.js
@@ -1,24 +1,28 @@
 const { web3 } = require('./setup');
 const { promisify } = require('util');
 
+const makeSnapshot = async () =>
+  await promisify(web3.currentProvider.send.bind(web3.currentProvider))({
+    jsonrpc: '2.0',
+    method: 'evm_snapshot',
+    id: new Date().getTime(),
+  }).then(d => d.result);
+
 /**
  * Returns a snapshot object with the 'restore' function, which reverts blockchain to the captured state
  */
 const snapshot = async function () {
-  const snapshotData = await promisify(web3.currentProvider.send.bind(web3.currentProvider))({
-    jsonrpc: '2.0',
-    method: 'evm_snapshot',
-    id: new Date().getTime(),
-  });
+  let snapshotId = await makeSnapshot();
 
   return {
     restore: async function () {
       await promisify(web3.currentProvider.send.bind(web3.currentProvider))({
         jsonrpc: '2.0',
         method: 'evm_revert',
-        params: [snapshotData.result],
+        params: [snapshotId],
         id: new Date().getTime(),
       });
+      snapshotId = await makeSnapshot();
     },
   };
 };

diff --git a/test/src/snapshot.test.js b/test/src/snapshot.test.js
@@ -28,10 +28,22 @@ describe('snapshot', function () {
   });
 
   contract('reverting transactions', function ([sender, receiver]) {
-    it('returns a snapshot object with restore method restoring previous ether balances', async function () {
+    it('restores previous balances on revert', async function () {
       const snapshotA = await snapshot();
       const tracker = await balance.tracker(receiver);
-      await tracker.get();
+      await send.ether(sender, receiver, ether('1'));
+      expect(await tracker.delta()).to.be.bignumber.equal(ether('1'));
+      await snapshotA.restore();
+      expect(await tracker.delta()).to.be.bignumber.equal(ether('-1'));
+    });
+
+    it('can revert twice', async function () {
+      const snapshotA = await snapshot();
+      const tracker = await balance.tracker(receiver);
+      await send.ether(sender, receiver, ether('1'));
+      expect(await tracker.delta()).to.be.bignumber.equal(ether('1'));
+      await snapshotA.restore();
+      expect(await tracker.delta()).to.be.bignumber.equal(ether('-1'));
       await send.ether(sender, receiver, ether('1'));
       expect(await tracker.delta()).to.be.bignumber.equal(ether('1'));
       await snapshotA.restore();