Use hardhat built-in SnapshotRestorer to handle snapshots

keep-network · Jan 16, 2024 · e089c95 · e089c95
1 parent 647ceed
commit e089c95
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 22 deletions.
diff --git a/package.json b/package.json
@@ -23,6 +23,7 @@
   "devDependencies": {
     "@keep-network/prettier-config-keep": "github:keep-network/prettier-config-keep",
     "@nomicfoundation/hardhat-ethers": "^3.0.5",
+    "@nomicfoundation/hardhat-network-helpers": "^1.0.10",
     "@nomicfoundation/hardhat-verify": "^2.0.3",
     "@openzeppelin/hardhat-upgrades": "^3.0.2",
     "@types/chai": "^4.2.21",

diff --git a/src/index.ts b/src/index.ts
@@ -59,7 +59,7 @@ extendEnvironment((hre) => {
         return signers(hre)
       }),
       snapshot: lazyObject(() => {
-        return snapshot(hre)
+        return snapshot()
       }),
       upgrades: lazyObject(() => {
         return upgrades(hre)

diff --git a/src/snapshot.ts b/src/snapshot.ts
@@ -1,7 +1,9 @@
-import { JsonRpcProvider } from "@ethersproject/providers"
-import { HardhatRuntimeEnvironment } from "hardhat/types"
+import {
+  SnapshotRestorer,
+  takeSnapshot,
+} from "@nomicfoundation/hardhat-network-helpers"
 
-const snapshotIdsStack: any[] = []
+const snapshotStack: SnapshotRestorer[] = []
 
 export interface HardhatSnapshotHelpers {
   createSnapshot(): Promise<void>
@@ -10,31 +12,23 @@ export interface HardhatSnapshotHelpers {
 
 /**
  * Snapshot the state of the blockchain at the current.
- * @param {JsonRpcProvider} provider Ethers provider
  */
-export async function createSnapshot(provider: JsonRpcProvider): Promise<void> {
-  const snapshotId = await provider.send("evm_snapshot", [])
-  snapshotIdsStack.push(snapshotId)
+export async function createSnapshot(): Promise<void> {
+  const snapshot = await takeSnapshot()
+  snapshotStack.push(snapshot)
 }
 
 /**
  * Restores the chain to a latest snapshot.
- * @param {JsonRpcProvider} provider Ethers provider
  */
-export async function restoreSnapshot(
-  provider: JsonRpcProvider
-): Promise<void> {
-  const snapshotId = snapshotIdsStack.pop()
-  await provider.send("evm_revert", [snapshotId])
+export async function restoreSnapshot(): Promise<void> {
+  const snapshot = snapshotStack.pop()
+  await snapshot?.restore()
 }
 
-export default function (
-  hre: HardhatRuntimeEnvironment
-): HardhatSnapshotHelpers {
-  const provider = hre.ethers.provider as unknown as JsonRpcProvider
-
+export default function (): HardhatSnapshotHelpers {
   return {
-    createSnapshot: () => createSnapshot(provider),
-    restoreSnapshot: () => restoreSnapshot(provider),
+    createSnapshot,
+    restoreSnapshot,
   }
 }
diff --git a/yarn.lock b/yarn.lock
@@ -885,6 +885,13 @@
     debug "^4.1.1"
     lodash.isequal "^4.5.0"
 
+"@nomicfoundation/hardhat-network-helpers@^1.0.10":
+  version "1.0.10"
+  resolved "https://registry.yarnpkg.com/@nomicfoundation/hardhat-network-helpers/-/hardhat-network-helpers-1.0.10.tgz#c61042ceb104fdd6c10017859fdef6529c1d6585"
+  integrity sha512-R35/BMBlx7tWN5V6d/8/19QCwEmIdbnA4ZrsuXgvs8i2qFx5i7h6mH5pBS4Pwi4WigLH+upl6faYusrNPuzMrQ==
+  dependencies:
+    ethereumjs-util "^7.1.4"
+
 "@nomicfoundation/hardhat-verify@^2.0.3":
   version "2.0.3"
   resolved "https://registry.yarnpkg.com/@nomicfoundation/hardhat-verify/-/hardhat-verify-2.0.3.tgz#173557f8cfa53c8c9da23a326f54d24fe459ae68"
@@ -2406,7 +2413,7 @@ ethereumjs-util@^7.0.3:
     ethereum-cryptography "^0.1.3"
     rlp "^2.2.4"
 
-ethereumjs-util@^7.1.5:
+ethereumjs-util@^7.1.4, ethereumjs-util@^7.1.5:
   version "7.1.5"
   resolved "https://registry.yarnpkg.com/ethereumjs-util/-/ethereumjs-util-7.1.5.tgz#9ecf04861e4fbbeed7465ece5f23317ad1129181"
   integrity sha512-SDl5kKrQAudFBUe5OJM9Ac6WmMyYmXX/6sTmLZ3ffG2eY6ZIGBes3pEDxNN6V72WyOw4CPD5RomKdsa8DAAwLg==