refactor(wrap/unwrap): txCallback is async

lidofinance · Nov 21, 2024 · 54da346 · 54da346
1 parent f0aea5b
commit 54da346
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 37 deletions.
diff --git a/features/wsteth/unwrap/hooks/use-unwrap-form-processing.ts b/features/wsteth/unwrap/hooks/use-unwrap-form-processing.ts
@@ -36,13 +36,6 @@ export const useUnwrapFormProcessor = ({
     processApproveTx: processApproveTxOnL2,
   } = approvalDataOnL2;
 
-  const showSuccessTxModal = useCallback(async () => {
-    const wstethBalance = await (isDappActiveOnL2
-      ? l2.steth.balance(address)
-      : stETH.balance(address));
-    txModalStages.success(wstethBalance, txHashRef.current);
-  }, [address, isDappActiveOnL2, l2.steth, stETH, txModalStages]);
-
   return useCallback(
     async ({ amount }: UnwrapFormInputType) => {
       try {
@@ -61,18 +54,24 @@ export const useUnwrapFormProcessor = ({
           // The operation 'wstETH to stETH' on L2 is 'wrap'
           await l2.wrapWstethToSteth({
             value: amount,
-            callback: ({ stage, payload }) => {
+            callback: async ({ stage, payload }) => {
               switch (stage) {
                 case TransactionCallbackStage.SIGN:
                   txModalStages.sign(amount, willReceive);
                   break;
                 case TransactionCallbackStage.RECEIPT:
+                  // the payload here is txHash
                   txModalStages.pending(amount, willReceive, payload);
-                  txHashRef.current = payload; // the payload here is txHash
+                  txHashRef.current = payload;
                   break;
                 case TransactionCallbackStage.DONE:
-                  void onConfirm?.();
-                  void showSuccessTxModal();
+                  await onConfirm?.();
+                  txModalStages.success(
+                    await (isDappActiveOnL2
+                      ? l2.steth.balance(address)
+                      : stETH.balance(address)),
+                    txHashRef.current,
+                  );
                   break;
                 case TransactionCallbackStage.MULTISIG_DONE:
                   txModalStages.successMultisig();
@@ -87,7 +86,7 @@ export const useUnwrapFormProcessor = ({
         } else {
           await wrap.unwrap({
             value: amount,
-            callback: ({ stage, payload }) => {
+            callback: async ({ stage, payload }) => {
               switch (stage) {
                 case TransactionCallbackStage.SIGN:
                   txModalStages.sign(amount, willReceive);
@@ -97,8 +96,13 @@ export const useUnwrapFormProcessor = ({
                   txHashRef.current = payload; // the payload here is txHash
                   break;
                 case TransactionCallbackStage.DONE:
-                  void onConfirm?.();
-                  void showSuccessTxModal();
+                  await onConfirm?.();
+                  txModalStages.success(
+                    await (isDappActiveOnL2
+                      ? l2.steth.balance(address)
+                      : stETH.balance(address)),
+                    txHashRef.current,
+                  );
                   break;
                 case TransactionCallbackStage.MULTISIG_DONE:
                   txModalStages.successMultisig();
@@ -130,7 +134,7 @@ export const useUnwrapFormProcessor = ({
       onRetry,
       txModalStages,
       onConfirm,
-      showSuccessTxModal,
+      stETH,
       wrap,
     ],
   );

diff --git a/features/wsteth/unwrap/hooks/use-unwrap-tx-on-l2-approve.ts b/features/wsteth/unwrap/hooks/use-unwrap-tx-on-l2-approve.ts
@@ -37,7 +37,7 @@ export const useUnwrapTxOnL2Approve = ({ amount }: UseUnwrapTxApproveArgs) => {
     async ({ onRetry }: { onRetry?: () => void }) => {
       const approveTxHash = await l2.approveWstethForWrap({
         value: amount,
-        callback: ({ stage, payload }) => {
+        callback: async ({ stage, payload }) => {
           switch (stage) {
             case TransactionCallbackStage.SIGN:
               txModalStages.signApproval(amount, TOKENS_TO_WRAP.ETH);

diff --git a/features/wsteth/wrap/hooks/use-wrap-form-processing.ts b/features/wsteth/wrap/hooks/use-wrap-form-processing.ts
@@ -38,13 +38,6 @@ export const useWrapFormProcessor = ({
     processApproveTx: processApproveTxOnL1,
   } = approvalDataOnL1;
 
-  const showSuccessTxModal = useCallback(async () => {
-    const wstethBalance = await (isDappActiveOnL2
-      ? l2.wsteth.balance(address)
-      : wstETH.balance(address));
-    txModalStages.success(wstethBalance, txHashRef.current);
-  }, [address, isDappActiveOnL2, l2.wsteth, txModalStages, wstETH]);
-
   return useCallback(
     async ({ amount, token }: WrapFormInputType) => {
       try {
@@ -59,18 +52,24 @@ export const useWrapFormProcessor = ({
           // The operation 'stETH to wstETH' on L2 is 'unwrap'
           await l2.unwrapStethToWsteth({
             value: amount,
-            callback: ({ stage, payload }) => {
+            callback: async ({ stage, payload }) => {
               switch (stage) {
                 case TransactionCallbackStage.SIGN:
                   txModalStages.sign(amount, token, willReceive);
                   break;
                 case TransactionCallbackStage.RECEIPT:
                   txModalStages.pending(amount, token, willReceive, payload);
-                  txHashRef.current = payload; // the payload here is txHash
+                  // the payload here is txHash
+                  txHashRef.current = payload;
                   break;
                 case TransactionCallbackStage.DONE:
-                  void onConfirm?.();
-                  void showSuccessTxModal();
+                  await onConfirm?.();
+                  txModalStages.success(
+                    await (isDappActiveOnL2
+                      ? l2.wsteth.balance(address)
+                      : wstETH.balance(address)),
+                    txHashRef.current,
+                  );
                   break;
                 case TransactionCallbackStage.MULTISIG_DONE:
                   txModalStages.successMultisig();
@@ -93,18 +92,24 @@ export const useWrapFormProcessor = ({
 
           await wrap.wrapSteth({
             value: amount,
-            callback: ({ stage, payload }) => {
+            callback: async ({ stage, payload }) => {
               switch (stage) {
                 case TransactionCallbackStage.SIGN:
                   txModalStages.sign(amount, token, willReceive);
                   break;
                 case TransactionCallbackStage.RECEIPT:
                   txModalStages.pending(amount, token, willReceive, payload);
-                  txHashRef.current = payload; // the payload here is txHash
+                  // the payload here is txHash
+                  txHashRef.current = payload;
                   break;
                 case TransactionCallbackStage.DONE:
-                  void onConfirm?.();
-                  void showSuccessTxModal();
+                  await onConfirm?.();
+                  txModalStages.success(
+                    await (isDappActiveOnL2
+                      ? l2.wsteth.balance(address)
+                      : wstETH.balance(address)),
+                    txHashRef.current,
+                  );
                   break;
                 case TransactionCallbackStage.MULTISIG_DONE:
                   txModalStages.successMultisig();
@@ -126,18 +131,24 @@ export const useWrapFormProcessor = ({
 
           await wrap.wrapEth({
             value: amount,
-            callback: ({ stage, payload }) => {
+            callback: async ({ stage, payload }) => {
               switch (stage) {
                 case TransactionCallbackStage.SIGN:
                   txModalStages.sign(amount, token, willReceive);
                   break;
                 case TransactionCallbackStage.RECEIPT:
                   txModalStages.pending(amount, token, willReceive, payload);
-                  txHashRef.current = payload; // the payload here is txHash
+                  // the payload here is txHash
+                  txHashRef.current = payload;
                   break;
                 case TransactionCallbackStage.DONE:
-                  void onConfirm?.();
-                  void showSuccessTxModal();
+                  await onConfirm?.();
+                  txModalStages.success(
+                    await (isDappActiveOnL2
+                      ? l2.wsteth.balance(address)
+                      : wstETH.balance(address)),
+                    txHashRef.current,
+                  );
                   break;
                 case TransactionCallbackStage.MULTISIG_DONE:
                   txModalStages.successMultisig();
@@ -166,7 +177,7 @@ export const useWrapFormProcessor = ({
       shares,
       txModalStages,
       onConfirm,
-      showSuccessTxModal,
+      wstETH,
       onRetry,
       isApprovalNeededBeforeWrapOnL1,
       wrap,

diff --git a/features/wsteth/wrap/hooks/use-wrap-tx-on-l1-approve.ts b/features/wsteth/wrap/hooks/use-wrap-tx-on-l1-approve.ts
@@ -50,7 +50,7 @@ export const useWrapTxOnL1Approve = ({
     async ({ onRetry }: { onRetry?: () => void }) => {
       const approveTx = await wrap.approveStethForWrap({
         value: amount,
-        callback: ({ stage, payload }) => {
+        callback: async ({ stage, payload }) => {
           switch (stage) {
             case TransactionCallbackStage.SIGN:
               txModalStages.signApproval(amount, token);