fix: applied suggestions

Author: vinistevam

packages/transaction-controller/CHANGELOG.md

@@ -12,7 +12,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add optional approval request when calling `addTransactionBatch` ([#5793](https://github.com/MetaMask/core/pull/5793))
   - Add `transactionBatches` array to state.
   - Add `TransactionBatchMeta` type.
-  - Add `addBatchMetadata` to store batch metadata and `wipeTransactionBatches` to clean up state after batch hook completion.
 
 ## [56.2.0]
 

packages/transaction-controller/src/utils/batch.test.ts

@@ -1402,7 +1402,7 @@ describe('Batch Utils', () => {
         expect(result?.batchId).toMatch(/^0x[0-9a-f]{32}$/u);
       });
 
-      it('saves a transaction batch and then cleans the state', async () => {
+      it('saves a transaction batch and then cleans the specific batch by ID', async () => {
         const { approve } = mockRequestApproval(MESSENGER_MOCK, {
           state: 'approved',
         });
@@ -1441,13 +1441,18 @@ describe('Batch Utils', () => {
         expect(updateMock).toHaveBeenCalledTimes(2);
         expect(updateMock).toHaveBeenCalledWith(expect.any(Function));
 
-        // Simulate the state update
+        // Simulate the state update for adding the batch
         const state = {
-          transactionBatches: [],
+          transactionBatches: [
+            { id: 'batch1', chainId: '0x1', transactions: [] },
+          ],
         } as unknown as TransactionControllerState;
+
+        // Simulate adding the batch
         updateMock.mock.calls[0][0](state);
 
         expect(state.transactionBatches).toStrictEqual([
+          { id: 'batch1', chainId: '0x1', transactions: [] },
           expect.objectContaining({
             id: expect.any(String),
             chainId: CHAIN_ID_MOCK,
@@ -1474,9 +1479,12 @@ describe('Batch Utils', () => {
 
         await resultPromise;
 
+        // Simulate cleaning the specific batch by ID
         updateMock.mock.calls[1][0](state);
 
-        expect(state.transactionBatches).toStrictEqual([]);
+        expect(state.transactionBatches).toStrictEqual([
+          { id: 'batch1', chainId: '0x1', transactions: [] },
+        ]);
       });
     });
   });

packages/transaction-controller/src/utils/batch.ts

@@ -54,7 +54,7 @@ import {
   TransactionType,
 } from '../types';
 
-type UpdateBatchMetadata = (
+type UpdateStateCallback = (
   callback: (
     state: WritableDraft<TransactionControllerState>,
   ) => void | TransactionControllerState,
@@ -81,7 +81,7 @@ type AddTransactionBatchRequest = {
   getPendingTransactionTracker: (
     networkClientId: string,
   ) => PendingTransactionTracker;
-  update: UpdateBatchMetadata;
+  update: UpdateStateCallback;
 };
 
 type IsAtomicBatchSupportedRequestInternal = {
@@ -480,8 +480,8 @@ async function addTransactionBatchWithHook(
 
     throw error;
   } finally {
-    log('Cleaning up publish batch hook');
-    wipeTransactionBatches(update);
+    log('Cleaning up publish batch hook', batchId);
+    wipeTransactionBatchById(update, batchId);
   }
 }
 
@@ -639,9 +639,9 @@ function newBatchMetadata({
  */
 function addBatchMetadata(
   transactionBatchMeta: TransactionBatchMeta,
-  update: UpdateBatchMetadata,
+  update: UpdateStateCallback,
 ) {
-  update((state: WritableDraft<TransactionControllerState>) => {
+  update((state) => {
     state.transactionBatches = [
       ...state.transactionBatches,
       transactionBatchMeta,
@@ -650,12 +650,18 @@ function addBatchMetadata(
 }
 
 /**
- * Wipes all transaction batches from the transaction controller state.
+ * Wipes a specific transaction batch from the transaction controller state by its ID.
  *
  * @param update - The update function to modify the transaction controller state.
+ * @param id - The ID of the transaction batch to be wiped.
  */
-function wipeTransactionBatches(update: UpdateBatchMetadata): void {
-  update((state: WritableDraft<TransactionControllerState>) => {
-    state.transactionBatches = [];
+function wipeTransactionBatchById(
+  update: UpdateStateCallback,
+  id: string,
+): void {
+  update((state) => {
+    state.transactionBatches = state.transactionBatches.filter(
+      (batch) => batch.id !== id,
+    );
   });
 }