From 40dad6e62d349b528304c5abdc3153658d4295cf Mon Sep 17 00:00:00 2001
From: Stanley Yuen <102275989+stanleyyconsensys@users.noreply.github.com>
Date: Wed, 9 Oct 2024 18:36:08 +0800
Subject: [PATCH] chore: change state mgr to return in withTransaction (#373)

---
 packages/starknet-snap/src/utils/snap-state.ts | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/packages/starknet-snap/src/utils/snap-state.ts b/packages/starknet-snap/src/utils/snap-state.ts
index 6902ba76..310bd345 100644
--- a/packages/starknet-snap/src/utils/snap-state.ts
+++ b/packages/starknet-snap/src/utils/snap-state.ts
@@ -67,10 +67,10 @@ export abstract class SnapStateManager<State> {
    *
    * @param callback - A Promise function that takes the state as an argument.
    */
-  public async withTransaction(
-    callback: (state: State) => Promise<void>,
-  ): Promise<void> {
-    await this.mtx.runExclusive(async () => {
+  public async withTransaction<Response>(
+    callback: (state: State) => Promise<Response>,
+  ): Promise<Response> {
+    return await this.mtx.runExclusive(async () => {
       await this.#beginTransaction();
 
       if (
@@ -88,8 +88,9 @@ export abstract class SnapStateManager<State> {
       );
 
       try {
-        await callback(this.#transaction.current);
+        const result = await callback(this.#transaction.current);
         await this.set(this.#transaction.current);
+        return result;
       } catch (error) {
         logger.info(
           `SnapStateManager.withTransaction [${