Split gas fee token request

Author: matthewwalsh0

packages/transaction-controller/src/TransactionController.test.ts

@@ -68,6 +68,7 @@ import type {
   InternalAccount,
   PublishHook,
   GasFeeToken,
+  SimulationData,
 } from './types';
 import {
   GasFeeEstimateType,
@@ -78,6 +79,7 @@ import {
   TransactionType,
   WalletDevice,
 } from './types';
+import { getBalanceChanges } from './utils/balance-changes';
 import { addTransactionBatch } from './utils/batch';
 import { DELEGATION_PREFIX, getDelegationAddress } from './utils/eip7702';
 import { addGasBuffer, estimateGas, updateGas } from './utils/gas';
@@ -87,8 +89,6 @@ import {
   getTransactionLayer1GasFee,
   updateTransactionLayer1GasFee,
 } from './utils/layer1-gas-fee-flow';
-import type { GetSimulationDataResult } from './utils/simulation';
-import { getSimulationData } from './utils/simulation';
 import {
   updatePostTransactionBalance,
   updateSwapsTransaction,
@@ -100,6 +100,7 @@ import {
   buildCustomNetworkClientConfiguration,
   buildMockGetNetworkClientById,
 } from '../../network-controller/tests/helpers';
+import { getGasFeeTokens } from './utils/gas-fee-tokens';
 
 type UnrestrictedMessenger = Messenger<
   TransactionControllerActions | AllowedActions,
@@ -128,8 +129,9 @@ jest.mock('./utils/batch');
 jest.mock('./utils/gas');
 jest.mock('./utils/gas-fees');
 jest.mock('./utils/gas-flow');
+jest.mock('./utils/gas-fee-tokens');
 jest.mock('./utils/layer1-gas-fee-flow');
-jest.mock('./utils/simulation');
+jest.mock('./utils/balance-changes');
 jest.mock('./utils/swaps');
 jest.mock('uuid');
 
@@ -453,27 +455,24 @@ const TRANSACTION_META_2_MOCK = {
   },
 } as TransactionMeta;
 
-const SIMULATION_DATA_RESULT_MOCK: GetSimulationDataResult = {
-  gasFeeTokens: [],
-  simulationData: {
-    nativeBalanceChange: {
-      previousBalance: '0x0',
-      newBalance: '0x1',
-      difference: '0x1',
+const SIMULATION_DATA_RESULT_MOCK: SimulationData = {
+  nativeBalanceChange: {
+    previousBalance: '0x0',
+    newBalance: '0x1',
+    difference: '0x1',
+    isDecrease: false,
+  },
+  tokenBalanceChanges: [
+    {
+      address: '0x123',
+      standard: SimulationTokenStandard.erc721,
+      id: '0x456',
+      previousBalance: '0x1',
+      newBalance: '0x3',
+      difference: '0x2',
       isDecrease: false,
     },
-    tokenBalanceChanges: [
-      {
-        address: '0x123',
-        standard: SimulationTokenStandard.erc721,
-        id: '0x456',
-        previousBalance: '0x1',
-        newBalance: '0x3',
-        difference: '0x2',
-        isDecrease: false,
-      },
-    ],
-  },
+  ],
 };
 
 const GAS_FEE_TOKEN_MOCK: GasFeeToken = {
@@ -523,7 +522,8 @@ describe('TransactionController', () => {
   );
   const testGasFeeFlowClassMock = jest.mocked(TestGasFeeFlow);
   const gasFeePollerClassMock = jest.mocked(GasFeePoller);
-  const getSimulationDataMock = jest.mocked(getSimulationData);
+  const getBalanceChangesMock = jest.mocked(getBalanceChanges);
+  const getGasFeeTokensMock = jest.mocked(getGasFeeTokens);
   const getTransactionLayer1GasFeeMock = jest.mocked(
     getTransactionLayer1GasFee,
   );
@@ -2168,7 +2168,7 @@ describe('TransactionController', () => {
 
     describe('updates simulation data', () => {
       it('by default', async () => {
-        getSimulationDataMock.mockResolvedValueOnce(
+        getBalanceChangesMock.mockResolvedValueOnce(
           SIMULATION_DATA_RESULT_MOCK,
         );
 
@@ -2186,27 +2186,23 @@ describe('TransactionController', () => {
 
         await flushPromises();
 
-        expect(getSimulationDataMock).toHaveBeenCalledTimes(1);
-        expect(getSimulationDataMock).toHaveBeenCalledWith(
-          {
-            chainId: MOCK_NETWORK.chainId,
-            data: undefined,
-            from: ACCOUNT_MOCK,
-            to: ACCOUNT_MOCK,
-            value: '0x0',
-          },
-          {
-            blockTime: undefined,
-          },
-        );
+        expect(getBalanceChangesMock).toHaveBeenCalledTimes(1);
+        expect(getBalanceChangesMock).toHaveBeenCalledWith({
+          blockTime: undefined,
+          chainId: MOCK_NETWORK.chainId,
+          data: undefined,
+          from: ACCOUNT_MOCK,
+          to: ACCOUNT_MOCK,
+          value: '0x0',
+        });
 
         expect(controller.state.transactions[0].simulationData).toStrictEqual(
-          SIMULATION_DATA_RESULT_MOCK.simulationData,
+          SIMULATION_DATA_RESULT_MOCK,
         );
       });
 
       it('with error if simulation disabled', async () => {
-        getSimulationDataMock.mockResolvedValueOnce(
+        getBalanceChangesMock.mockResolvedValueOnce(
           SIMULATION_DATA_RESULT_MOCK,
         );
 
@@ -2224,7 +2220,7 @@ describe('TransactionController', () => {
           },
         );
 
-        expect(getSimulationDataMock).toHaveBeenCalledTimes(0);
+        expect(getBalanceChangesMock).toHaveBeenCalledTimes(0);
         expect(controller.state.transactions[0].simulationData).toStrictEqual({
           error: {
             code: SimulationErrorCode.Disabled,
@@ -2235,7 +2231,7 @@ describe('TransactionController', () => {
       });
 
       it('unless approval not required', async () => {
-        getSimulationDataMock.mockResolvedValueOnce(
+        getBalanceChangesMock.mockResolvedValueOnce(
           SIMULATION_DATA_RESULT_MOCK,
         );
 
@@ -2249,12 +2245,12 @@ describe('TransactionController', () => {
           { requireApproval: false, networkClientId: NETWORK_CLIENT_ID_MOCK },
         );
 
-        expect(getSimulationDataMock).toHaveBeenCalledTimes(0);
+        expect(getBalanceChangesMock).toHaveBeenCalledTimes(0);
         expect(controller.state.transactions[0].simulationData).toBeUndefined();
       });
 
       it('with sender code if type 4', async () => {
-        getSimulationDataMock.mockResolvedValueOnce(
+        getBalanceChangesMock.mockResolvedValueOnce(
           SIMULATION_DATA_RESULT_MOCK,
         );
 
@@ -2277,20 +2273,17 @@ describe('TransactionController', () => {
 
         await flushPromises();
 
-        expect(getSimulationDataMock).toHaveBeenCalledWith(expect.any(Object), {
-          senderCode: DELEGATION_PREFIX + ACCOUNT_2_MOCK.slice(2),
-        });
+        expect(getBalanceChangesMock).toHaveBeenCalledWith(
+          expect.objectContaining({
+            senderCode: DELEGATION_PREFIX + ACCOUNT_2_MOCK.slice(2),
+          }),
+        );
       });
     });
 
     describe('updates gas fee tokens', () => {
       it('by default', async () => {
-        getSimulationDataMock.mockResolvedValueOnce({
-          gasFeeTokens: [GAS_FEE_TOKEN_MOCK],
-          simulationData: {
-            tokenBalanceChanges: [],
-          },
-        });
+        getGasFeeTokensMock.mockResolvedValueOnce([GAS_FEE_TOKEN_MOCK]);
 
         const { controller } = setupController();
 
@@ -2312,12 +2305,7 @@ describe('TransactionController', () => {
       });
 
       it('unless approval not required', async () => {
-        getSimulationDataMock.mockResolvedValueOnce({
-          gasFeeTokens: [GAS_FEE_TOKEN_MOCK],
-          simulationData: {
-            tokenBalanceChanges: [],
-          },
-        });
+        getGasFeeTokensMock.mockResolvedValueOnce([GAS_FEE_TOKEN_MOCK]);
 
         const { controller } = setupController();
 
@@ -2329,7 +2317,7 @@ describe('TransactionController', () => {
           { requireApproval: false, networkClientId: NETWORK_CLIENT_ID_MOCK },
         );
 
-        expect(getSimulationDataMock).toHaveBeenCalledTimes(0);
+        expect(getBalanceChangesMock).toHaveBeenCalledTimes(0);
         expect(controller.state.transactions[0].gasFeeTokens).toBeUndefined();
       });
     });
@@ -6837,7 +6825,7 @@ describe('TransactionController', () => {
         updateToInitialState: true,
       });
 
-      expect(getSimulationDataMock).toHaveBeenCalledTimes(0);
+      expect(getBalanceChangesMock).toHaveBeenCalledTimes(0);
 
       shouldResimulateMock.mockReturnValueOnce({
         blockTime: 123,
@@ -6848,17 +6836,13 @@ describe('TransactionController', () => {
 
       await flushPromises();
 
-      expect(getSimulationDataMock).toHaveBeenCalledTimes(1);
-      expect(getSimulationDataMock).toHaveBeenCalledWith(
-        {
-          from: ACCOUNT_MOCK,
-          to: ACCOUNT_2_MOCK,
-          value: TRANSACTION_META_MOCK.txParams.value,
-        },
-        {
-          blockTime: 123,
-        },
-      );
+      expect(getBalanceChangesMock).toHaveBeenCalledTimes(1);
+      expect(getBalanceChangesMock).toHaveBeenCalledWith({
+        blockTime: 123,
+        from: ACCOUNT_MOCK,
+        to: ACCOUNT_2_MOCK,
+        value: TRANSACTION_META_MOCK.txParams.value,
+      });
     });
 
     it('does not trigger simulation loop', async () => {
@@ -6876,7 +6860,7 @@ describe('TransactionController', () => {
         updateToInitialState: true,
       });
 
-      expect(getSimulationDataMock).toHaveBeenCalledTimes(0);
+      expect(getBalanceChangesMock).toHaveBeenCalledTimes(0);
 
       shouldResimulateMock.mockReturnValue({
         blockTime: 123,
@@ -6887,17 +6871,13 @@ describe('TransactionController', () => {
 
       await flushPromises();
 
-      expect(getSimulationDataMock).toHaveBeenCalledTimes(1);
-      expect(getSimulationDataMock).toHaveBeenCalledWith(
-        {
-          from: ACCOUNT_MOCK,
-          to: ACCOUNT_2_MOCK,
-          value: TRANSACTION_META_MOCK.txParams.value,
-        },
-        {
-          blockTime: 123,
-        },
-      );
+      expect(getBalanceChangesMock).toHaveBeenCalledTimes(1);
+      expect(getBalanceChangesMock).toHaveBeenCalledWith({
+        blockTime: 123,
+        from: ACCOUNT_MOCK,
+        to: ACCOUNT_2_MOCK,
+        value: TRANSACTION_META_MOCK.txParams.value,
+      });
     });
   });
 

packages/transaction-controller/src/TransactionController.ts

@@ -118,6 +118,7 @@ import {
   TransactionStatus,
   SimulationErrorCode,
 } from './types';
+import { getBalanceChanges } from './utils/balance-changes';
 import { addTransactionBatch, isAtomicBatchSupported } from './utils/batch';
 import {
   DELEGATION_PREFIX,
@@ -127,6 +128,7 @@ import {
 } from './utils/eip7702';
 import { validateConfirmedExternalTransaction } from './utils/external-transactions';
 import { addGasBuffer, estimateGas, updateGas } from './utils/gas';
+import { getGasFeeTokens } from './utils/gas-fee-tokens';
 import { updateGasFees } from './utils/gas-fees';
 import { getGasFeeFlow } from './utils/gas-flow';
 import {
@@ -143,7 +145,6 @@ import {
 } from './utils/nonce';
 import { prepareTransaction, serializeTransaction } from './utils/prepare';
 import { getTransactionParamsWithIncreasedGasFee } from './utils/retry';
-import { getSimulationData } from './utils/simulation';
 import {
   updatePostTransactionBalance,
   updateSwapsTransaction,
@@ -3986,27 +3987,20 @@ export class TransactionController extends BaseController<
         authorizationAddress &&
         ((DELEGATION_PREFIX + remove0x(authorizationAddress)) as Hex);
 
-      const result = await this.#trace(
+      simulationData = await this.#trace(
         { name: 'Simulate', parentContext: traceContext },
         () =>
-          getSimulationData(
-            {
-              chainId,
-              from: from as Hex,
-              to: to as Hex,
-              value: value as Hex,
-              data: data as Hex,
-            },
-            {
-              blockTime,
-              senderCode,
-            },
-          ),
+          getBalanceChanges({
+            blockTime,
+            chainId,
+            data: data as Hex,
+            from: from as Hex,
+            senderCode,
+            to: to as Hex,
+            value: value as Hex,
+          }),
       );
 
-      gasFeeTokens = result?.gasFeeTokens;
-      simulationData = result?.simulationData;
-
       if (
         blockTime &&
         prevSimulationData &&
@@ -4017,19 +4011,11 @@ export class TransactionController extends BaseController<
           isUpdatedAfterSecurityCheck: true,
         };
       }
-    }
 
-    const finalTransactionMeta = this.#getTransaction(transactionId);
-
-    /* istanbul ignore if */
-    if (!finalTransactionMeta) {
-      log(
-        'Cannot update simulation data as transaction not found',
-        transactionId,
-        simulationData,
-      );
-
-      return;
+      gasFeeTokens = await getGasFeeTokens({
+        chainId,
+        txParams,
+      });
     }
 
     this.#updateTransactionInternal(