Use new assets state in TransactionPayController (behind feature flag)

packages/transaction-pay-controller/CHANGELOG.md

@@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- **BREAKING:** Add `AssetsControllerGetStateForTransactionPayAction` to the `AllowedActions` messenger type ([#8163](https://github.com/MetaMask/core/pull/8163))
+
+### Changed
+
+- `getTokenBalance`, `getTokenInfo`, and `getTokenFiatRate` now source token metadata, balances, and pricing from `AssetsControllerGetStateForTransactionPayAction` when the `assetsUnifyState` remote feature flag is enabled, falling back to individual controller state calls otherwise ([#8163](https://github.com/MetaMask/core/pull/8163))
+
 ## [16.5.0]
 
 ### Added

packages/transaction-pay-controller/package.json

@@ -51,6 +51,7 @@
     "@ethersproject/abi": "^5.7.0",
     "@ethersproject/contracts": "^5.7.0",
     "@ethersproject/providers": "^5.7.0",
+    "@metamask/assets-controller": "^2.3.0",
     "@metamask/assets-controllers": "^100.2.1",
     "@metamask/base-controller": "^9.0.0",
     "@metamask/bridge-controller": "^69.1.0",

packages/transaction-pay-controller/src/tests/messenger-mock.ts

@@ -127,6 +127,8 @@ export function getMessengerMock({
     TransactionControllerEstimateGasBatchAction['handler']
   > = jest.fn();
 
+  const getAssetsControllerStateMock = jest.fn();
+
   const messenger: RootMessenger = new Messenger({
     namespace: MOCK_ANY_NAMESPACE,
   });
@@ -241,12 +243,18 @@ export function getMessengerMock({
       'TransactionController:estimateGasBatch',
       estimateGasBatchMock,
     );
+
+    messenger.registerActionHandler(
+      'AssetsController:getStateForTransactionPay',
+      getAssetsControllerStateMock,
+    );
   }
 
   const publish = messenger.publish.bind(messenger);
 
   return {
     addTransactionMock,
+    getAssetsControllerStateMock,
     addTransactionBatchMock,
     estimateGasMock,
     estimateGasBatchMock,

packages/transaction-pay-controller/src/types.ts

@@ -1,3 +1,4 @@
+import type { AssetsControllerGetStateForTransactionPayAction } from '@metamask/assets-controller';
 import type {
   CurrencyRateControllerActions,
   TokenBalancesControllerGetStateAction,
@@ -38,6 +39,7 @@ import type { CONTROLLER_NAME, TransactionPayStrategy } from './constants';
 
 export type AllowedActions =
   | AccountTrackerControllerGetStateAction
+  | AssetsControllerGetStateForTransactionPayAction
   | BridgeControllerActions
   | BridgeStatusControllerActions
   | CurrencyRateControllerActions

packages/transaction-pay-controller/src/utils/feature-flags.test.ts

@@ -9,6 +9,7 @@ import {
   DEFAULT_RELAY_QUOTE_URL,
   DEFAULT_SLIPPAGE,
   DEFAULT_STRATEGY_ORDER,
+  getAssetsUnifyStateFeature,
   getFallbackGas,
   DEFAULT_RELAY_EXECUTE_URL,
   getRelayOriginGasOverhead,
@@ -562,6 +563,85 @@ describe('Feature Flags Utils', () => {
     });
   });
 
+  describe('getAssetsUnifyStateFeature', () => {
+    type AssetsUnifyingState =
+      | {
+          enabled: boolean;
+          featureVersion: string | null;
+        }
+      | undefined;
+
+    const failureCases: {
+      description: string;
+      assetsUnifyingState: AssetsUnifyingState;
+    }[] = [
+      {
+        description: 'returns false when assetsUnifyState is not set',
+        assetsUnifyingState: undefined,
+      },
+      {
+        description: 'returns false when assetsUnifyState.enabled is false',
+        assetsUnifyingState: {
+          enabled: false,
+          featureVersion: '1',
+        },
+      },
+      {
+        description:
+          'returns false when featureVersion does not match expected version',
+        assetsUnifyingState: {
+          enabled: true,
+          featureVersion: '2',
+        },
+      },
+    ];
+
+    const successCases = [
+      {
+        description:
+          'returns true when assetsUnifyState is enabled and featureVersion matches',
+        assetsUnifyingState: {
+          enabled: true,
+          featureVersion: '1',
+        },
+      },
+    ];
+
+    const arrangeMocks = (assetsUnifyState: AssetsUnifyingState): void => {
+      const defaultRemoteFeatureFlagsState =
+        getDefaultRemoteFeatureFlagControllerState();
+      getRemoteFeatureFlagControllerStateMock.mockReturnValue({
+        ...defaultRemoteFeatureFlagsState,
+        remoteFeatureFlags: {
+          ...defaultRemoteFeatureFlagsState.remoteFeatureFlags,
+          ...(assetsUnifyState ? { assetsUnifyState } : {}),
+        },
+      });
+    };
+
+    it.each(failureCases)(
+      '$description',
+      ({ assetsUnifyingState }: (typeof failureCases)[number]) => {
+        arrangeMocks(assetsUnifyingState);
+
+        const result = getAssetsUnifyStateFeature(messenger);
+
+        expect(result).toBe(false);
+      },
+    );
+
+    it.each(successCases)(
+      '$description',
+      ({ assetsUnifyingState }: (typeof successCases)[number]) => {
+        arrangeMocks(assetsUnifyingState);
+
+        const result = getAssetsUnifyStateFeature(messenger);
+
+        expect(result).toBe(true);
+      },
+    );
+  });
+
   describe('getStrategyOrder', () => {
     it('returns default strategy order when none is set', () => {
       const strategyOrder = getStrategyOrder(messenger);

packages/transaction-pay-controller/src/utils/feature-flags.ts

@@ -297,6 +297,31 @@ export function getSlippage(
   return slippage;
 }
 
+/**
+ * Get the AssetsUnifyState feature flag state.
+ *
+ * @param messenger - Controller messenger.
+ * @returns True if the assets unify state feature is enabled, false otherwise.
+ */
+export function getAssetsUnifyStateFeature(
+  messenger: TransactionPayControllerMessenger,
+): boolean {
+  const state = messenger.call('RemoteFeatureFlagController:getState');
+  const assetsUnifyState = state.remoteFeatureFlags.assetsUnifyState as
+    | {
+        enabled: boolean;
+        featureVersion: string | null;
+      }
+    | undefined;
+
+  const AssetsUnifyStateFeatureVersion = '1';
+
+  return (
+    Boolean(assetsUnifyState?.enabled) &&
+    assetsUnifyState?.featureVersion === AssetsUnifyStateFeatureVersion
+  );
+}
+
 /**
  * Get a value from a record using a case-insensitive key lookup.
  *