diff --git a/src/lib/utils.ts b/src/lib/utils.ts
index 05580cfd..0a564ae4 100644
--- a/src/lib/utils.ts
+++ b/src/lib/utils.ts
@@ -75,13 +75,13 @@ let snapshotHeight = -1;
 
 export const createSnapshot = async (aeSdk) => {
   snapshotHeight = await aeSdk.getHeight();
-  await awaitKeyBlocks(aeSdk, 1);
 };
 
-export const rollbackSnapshot = async (aeSdk) => {
+export const rollbackHeight = async (aeSdk, height) => {
   const currentBlockHeight = await aeSdk.getHeight();
-  if (currentBlockHeight > snapshotHeight) {
-    await get(`http://localhost:3001/rollback?height=${snapshotHeight}`);
-    await awaitKeyBlocks(aeSdk, 1);
+  if (currentBlockHeight > height) {
+    await get(`http://localhost:3001/rollback?height=${height}`);
   }
 };
+
+export const rollbackSnapshot = async (aeSdk) => rollbackHeight(aeSdk, snapshotHeight);