fix(predict): address confirmation hook review

Author: matallui

app/components/UI/Predict/controllers/PredictController.ts

@@ -2732,7 +2732,7 @@ export class PredictController extends BaseController<
 
   public async publish(_request: {
     transactionMeta: TransactionMeta;
-  }): Promise<{ transactionHash?: string; isIntentComplete?: boolean }> {
+  }): Promise<{ transactionHash?: string }> {
     return { transactionHash: undefined };
   }
 

app/core/Engine/controllers/transaction-controller/transaction-controller-init.test.ts

@@ -122,22 +122,43 @@ function buildInitRequestMock(
     TransactionControllerInitMessenger
   >
 > {
+  const {
+    predictControllerMock: providedPredictControllerMock,
+    ...requestOverrides
+  } = initRequestProperties;
+  const predictControllerMock =
+    (providedPredictControllerMock as ControllerMock | undefined) ??
+    buildControllerMock();
   const initMessenger = new ExtendedMessenger<MockAnyNamespace>({
     namespace: MOCK_ANY_NAMESPACE,
   });
   const baseControllerMessenger = new ExtendedMessenger<MockAnyNamespace>({
     namespace: MOCK_ANY_NAMESPACE,
   });
+  (initMessenger as unknown as { call: jest.Mock }).call = jest.fn(
+    (actionType: string, params: unknown) => {
+      if (actionType === 'PredictController:beforePublish') {
+        return predictControllerMock.beforePublish(params);
+      }
+
+      if (actionType === 'PredictController:publish') {
+        return predictControllerMock.publish(params);
+      }
+
+      throw new Error(`Unexpected init messenger action: ${actionType}`);
+    },
+  );
+
   const requestMock = {
     ...buildMessengerClientInitRequestMock(baseControllerMessenger),
     initMessenger:
       initMessenger as unknown as TransactionControllerInitMessenger,
     controllerMessenger:
       baseControllerMessenger as unknown as TransactionControllerMessenger,
-    ...initRequestProperties,
+    ...requestOverrides,
   };
 
-  if (!initRequestProperties.getMessengerClient) {
+  if (!requestOverrides.getMessengerClient) {
     requestMock.getMessengerClient.mockReturnValue(buildControllerMock());
   }
 
@@ -339,11 +360,7 @@ describe('Transaction Controller Init', () => {
         'hooks',
         {},
         {
-          getMessengerClient: jest.fn((controllerName: string) =>
-            controllerName === 'PredictController'
-              ? predictControllerMock
-              : buildControllerMock(),
-          ),
+          predictControllerMock,
         },
       );
 
@@ -388,11 +405,7 @@ describe('Transaction Controller Init', () => {
         'hooks',
         {},
         {
-          getMessengerClient: jest.fn((controllerName: string) =>
-            controllerName === 'PredictController'
-              ? predictControllerMock
-              : buildControllerMock(),
-          ),
+          predictControllerMock,
         },
       );
 
@@ -422,11 +435,7 @@ describe('Transaction Controller Init', () => {
         'hooks',
         {},
         {
-          getMessengerClient: jest.fn((controllerName: string) =>
-            controllerName === 'PredictController'
-              ? predictControllerMock
-              : buildControllerMock(),
-          ),
+          predictControllerMock,
         },
       );
 
@@ -437,52 +446,6 @@ describe('Transaction Controller Init', () => {
       expect(submitSmartTransactionHookMock).not.toHaveBeenCalled();
     });
 
-    it('marks the latest transaction intent complete when Predict publish completes an intent', async () => {
-      const predictControllerMock = buildControllerMock({
-        publish: jest.fn().mockResolvedValue({
-          transactionHash: '0xpredict',
-          isIntentComplete: true,
-        }),
-      } as unknown as Partial<NetworkController>);
-      const getTransactionByIdMock = jest.requireMock(
-        '../../../../util/transactions',
-      ).getTransactionById;
-      getTransactionByIdMock.mockReturnValue({ ...MOCK_TRANSACTION_META });
-      const hooks = testConstructorOption(
-        'hooks',
-        {},
-        {
-          getMessengerClient: jest.fn((controllerName: string) =>
-            controllerName === 'PredictController'
-              ? predictControllerMock
-              : buildControllerMock(),
-          ),
-        },
-      );
-
-      const result = await hooks?.publish?.(MOCK_TRANSACTION_META);
-
-      const transactionControllerInstance = transactionControllerClassMock.mock
-        .instances[0] as unknown as {
-        updateTransaction: jest.Mock;
-      };
-
-      expect(result).toEqual({ transactionHash: '0xpredict' });
-      expect(getTransactionByIdMock).toHaveBeenCalledWith(
-        MOCK_TRANSACTION_META.id,
-        transactionControllerInstance,
-      );
-      expect(
-        transactionControllerInstance.updateTransaction,
-      ).toHaveBeenCalledWith(
-        expect.objectContaining({
-          id: MOCK_TRANSACTION_META.id,
-          isIntentComplete: true,
-        }),
-        'Predict claim relayer intent complete',
-      );
-    });
-
     it('passes isSmartTransaction returning false to pay hook when stxDisabled is true', async () => {
       selectMetaMaskPayFlagsMock.mockReturnValue({
         attemptsMax: 2,

app/core/Engine/controllers/transaction-controller/transaction-controller-init.ts

@@ -125,7 +125,6 @@ export const TransactionControllerInit: MessengerClientInitFunction<
               transactionController,
               smartTransactionsController,
               initMessenger,
-              request,
               signedTransactionInHex,
             }),
           publishBatch: async (_request: PublishBatchHookRequest) =>
@@ -138,7 +137,7 @@ export const TransactionControllerInit: MessengerClientInitFunction<
                 _request.transactions as PublishBatchHookTransaction[],
             }),
           beforePublish: (transactionMeta: TransactionMeta) =>
-            beforePublish(transactionMeta, request),
+            beforePublish(transactionMeta, initMessenger),
           beforeSign: (_request: { transactionMeta: TransactionMeta }) =>
             beforeSign(_request, request),
         },
@@ -219,7 +218,6 @@ async function publishHook({
   transactionController,
   smartTransactionsController,
   initMessenger,
-  request,
   signedTransactionInHex,
 }: {
   transactionMeta: TransactionMeta;
@@ -228,20 +226,15 @@ async function publishHook({
   transactionController: TransactionController;
   smartTransactionsController: SmartTransactionsController;
   initMessenger: TransactionControllerInitMessenger;
-  request: MessengerClientInitRequest<
-    TransactionControllerMessenger,
-    TransactionControllerInitMessenger
-  >;
   signedTransactionInHex: Hex;
 }): Promise<{ transactionHash?: string }> {
-  const predictResult = await publishPredict({
-    transactionMeta,
-    transactionController,
-    request,
-  });
+  const { transactionHash: predictTransactionHash } = await initMessenger.call(
+    'PredictController:publish',
+    { transactionMeta },
+  );
 
-  if (predictResult.transactionHash) {
-    return { transactionHash: predictResult.transactionHash };
+  if (predictTransactionHash) {
+    return { transactionHash: predictTransactionHash };
   }
 
   const state = getState();
@@ -344,41 +337,6 @@ async function publishHook({
   return { transactionHash: undefined };
 }
 
-async function publishPredict({
-  transactionMeta,
-  transactionController,
-  request,
-}: {
-  transactionMeta: TransactionMeta;
-  transactionController: TransactionController;
-  request: MessengerClientInitRequest<
-    TransactionControllerMessenger,
-    TransactionControllerInitMessenger
-  >;
-}): Promise<{ transactionHash?: string; isIntentComplete?: boolean }> {
-  const predictController = request.getMessengerClient('PredictController');
-  const result = await predictController.publish({ transactionMeta });
-
-  if (result.transactionHash && result.isIntentComplete) {
-    const latestMeta = getTransactionById(
-      transactionMeta.id,
-      transactionController,
-    );
-
-    if (latestMeta) {
-      transactionController.updateTransaction(
-        {
-          ...latestMeta,
-          isIntentComplete: true,
-        },
-        'Predict claim relayer intent complete',
-      );
-    }
-  }
-
-  return result;
-}
-
 function getSmartTransactionCommonParams(state: RootState, chainId: Hex) {
   const shouldUseSmartTransaction = selectShouldUseSmartTransaction(
     state,
@@ -496,13 +454,11 @@ function getControllers(
 
 function beforePublish(
   transactionMeta: TransactionMeta,
-  request: MessengerClientInitRequest<
-    TransactionControllerMessenger,
-    TransactionControllerInitMessenger
-  >,
+  initMessenger: TransactionControllerInitMessenger,
 ) {
-  const predictController = request.getMessengerClient('PredictController');
-  return predictController.beforePublish({ transactionMeta });
+  return initMessenger.call('PredictController:beforePublish', {
+    transactionMeta,
+  });
 }
 
 function beforeSign(

app/core/Engine/messengers/transaction-controller-messenger/transaction-controller-messenger.ts

@@ -55,6 +55,10 @@ import {
   MessengerActions,
   MessengerEvents,
 } from '@metamask/messenger';
+import type {
+  PredictControllerBeforePublishAction,
+  PredictControllerPublishAction,
+} from '../../../../components/UI/Predict/controllers/PredictController-method-action-types';
 
 export function getTransactionControllerMessenger(
   rootMessenger: RootMessenger,
@@ -114,7 +118,9 @@ type InitMessengerActions =
   | TransactionPayControllerGetDelegationTransactionAction
   | TransactionPayControllerGetStateAction
   | TransactionPayControllerGetStrategyAction
-  | AnalyticsControllerActions;
+  | AnalyticsControllerActions
+  | PredictControllerBeforePublishAction
+  | PredictControllerPublishAction;
 
 type InitMessengerEvents =
   | BridgeStatusControllerEvents
@@ -173,6 +179,8 @@ export function getTransactionControllerInitMessenger(
       'TransactionPayController:getState',
       'TransactionPayController:getStrategy',
       'AnalyticsController:trackEvent',
+      'PredictController:beforePublish',
+      'PredictController:publish',
     ],
     events: [
       'BridgeStatusController:stateChange',