Make it possible to have multiple layers of snapshots

1 files changed, 7 insertions, 4 deletions

diff --git a/test/utils/blockchain_lifecycle.ts b/test/utils/blockchain_lifecycle.ts
index 50eb57b95..9fdf0e856 100644
--- a/test/utils/blockchain_lifecycle.ts
+++ b/test/utils/blockchain_lifecycle.ts
@@ -2,19 +2,22 @@ import {RPC} from './rpc';
 
 export class BlockchainLifecycle {
     private rpc: RPC;
-    private snapshotId: number;
+    private snapshotIdsStack: number[];
     constructor() {
         this.rpc = new RPC();
+        this.snapshotIdsStack = [];
     }
     // TODO: In order to run these tests on an actual node, we should check if we are running against
     // TestRPC, if so, use snapshots, otherwise re-deploy contracts before every test
     public async startAsync(): Promise<void> {
-        this.snapshotId = await this.rpc.takeSnapshotAsync();
+        const snapshotId = await this.rpc.takeSnapshotAsync();
+        this.snapshotIdsStack.push(snapshotId);
     }
     public async revertAsync(): Promise<void> {
-        const didRevert = await this.rpc.revertSnapshotAsync(this.snapshotId);
+        const snapshotId = this.snapshotIdsStack.pop() as number;
+        const didRevert = await this.rpc.revertSnapshotAsync(snapshotId);
         if (!didRevert) {
-            throw new Error(`Snapshot with id #${this.snapshotId} failed to revert`);
+            throw new Error(`Snapshot with id #${snapshotId} failed to revert`);
         }
     }
 }