From 03098589172a49f413e390e823fd04a1e9f68d5c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Albert=20Oliv=C3=A9?= <albertolivecorbella@gmail.com>
Date: Mon, 17 Jul 2023 12:11:38 +0200
Subject: [PATCH] =?UTF-8?q?[MMI]=C2=A0Fixed=20remove=20custodian=20token?=
 =?UTF-8?q?=20(#20021)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 app/scripts/controllers/mmi-controller.js     |  21 +--
 .../__snapshots__/checkbox.test.tsx.snap      |   2 +-
 .../confirm-remove-jwt-modal.test.js.snap     | 154 +-----------------
 .../confirm-remove-jwt-modal.js               | 113 ++++++++-----
 .../confirm-remove-jwt-modal.stories.js       |   2 +-
 .../confirm-remove-jwt-modal.test.js          |  11 +-
 .../confirm-remove-jwt-modal/index.scss       |   1 +
 ...eractive-replacement-token-notification.js |   2 +-
 .../account-list-item-menu.js                 |   4 +-
 ui/pages/institutional/custody/custody.js     |   6 +-
 10 files changed, 86 insertions(+), 230 deletions(-)

diff --git a/app/scripts/controllers/mmi-controller.js b/app/scripts/controllers/mmi-controller.js
index 8923487b5..0ed4deaae 100644
--- a/app/scripts/controllers/mmi-controller.js
+++ b/app/scripts/controllers/mmi-controller.js
@@ -442,25 +442,8 @@ export default class MMIController extends EventEmitter {
     return keyring.getTransactionDeepLink(from, custodyTxId);
   }
 
-  async getCustodianToken(custodianType) {
-    let currentCustodyType;
-
-    const address = this.preferencesController.getSelectedAddress();
-
-    if (!custodianType) {
-      const resultCustody = this.custodyController.getCustodyTypeByAddress(
-        toChecksumHexAddress(address),
-      );
-      currentCustodyType = resultCustody;
-    }
-    let keyring = await this.keyringController.getKeyringsByType(
-      currentCustodyType || `Custody - ${custodianType}`,
-    )[0];
-    if (!keyring) {
-      keyring = await this.keyringController.addNewKeyring(
-        currentCustodyType || `Custody - ${custodianType}`,
-      );
-    }
+  async getCustodianToken(address) {
+    const keyring = await this.keyringController.getKeyringForAccount(address);
     const { authDetails } = keyring.getAccountDetails(address);
     return keyring ? authDetails.jwt || authDetails.refreshToken : '';
   }
diff --git a/ui/components/component-library/checkbox/__snapshots__/checkbox.test.tsx.snap b/ui/components/component-library/checkbox/__snapshots__/checkbox.test.tsx.snap
index 5df6da8de..90c0f27f5 100644
--- a/ui/components/component-library/checkbox/__snapshots__/checkbox.test.tsx.snap
+++ b/ui/components/component-library/checkbox/__snapshots__/checkbox.test.tsx.snap
@@ -3,7 +3,7 @@
 exports[`Checkbox should render the Checkbox without crashing 1`] = `
 <div>
   <label
-    class="box mm-text mm-checkbox mm-text--body-md box--display-inline-flex box--flex-direction-row box--align-items-center box--color-text-default"
+    class="mm-box mm-text mm-checkbox mm-text--body-md mm-box--display-inline-flex mm-box--align-items-center mm-box--color-text-default"
   >
     <span
       class="mm-checkbox__input-wrapper"
diff --git a/ui/components/institutional/confirm-remove-jwt-modal/__snapshots__/confirm-remove-jwt-modal.test.js.snap b/ui/components/institutional/confirm-remove-jwt-modal/__snapshots__/confirm-remove-jwt-modal.test.js.snap
index ccca0a8eb..7f057c9fe 100644
--- a/ui/components/institutional/confirm-remove-jwt-modal/__snapshots__/confirm-remove-jwt-modal.test.js.snap
+++ b/ui/components/institutional/confirm-remove-jwt-modal/__snapshots__/confirm-remove-jwt-modal.test.js.snap
@@ -1,155 +1,3 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
-exports[`Confirm Remove JWT should render correctly 1`] = `
-<div>
-  <div
-    class="modal-container"
-  >
-    <div
-      class="modal-container__header"
-    >
-      <div
-        class="modal-container__header-text"
-      >
-        Remove custodian token?
-      </div>
-      <div
-        class="modal-container__header-close"
-        data-testid="modal-header-close"
-      />
-    </div>
-    <div
-      class="modal-container__content"
-    >
-      <div
-        class="box confirm-action-jwt__jwt box--padding-2 box--display-flex box--flex-direction-row box--rounded-sm"
-      >
-        ...Dce07538D
-      </div>
-      <p
-        class="box mm-text confirm-action-jwt__show-more mm-text--body-md box--margin-left-2 box--flex-direction-row box--color-goerli"
-      >
-        <a
-          rel="noopener noreferrer"
-        >
-          Show more
-        </a>
-      </p>
-      <h6
-        class="box mm-text mm-text--body-sm mm-text--text-align-center box--margin-top-2 box--flex-direction-row box--color-text-default"
-      >
-        Are you sure you want to remove this token? All accounts assigned to this token will be removed from extension as well: 
-      </h6>
-      <div
-        class="box confirm-action-jwt__accounts-list box--flex-direction-row"
-      >
-        <div
-          class="mm-box page-container"
-        >
-          <div
-            class="mm-box page-container__content mm-box--padding-4"
-          >
-            <div
-              class="mm-box custody-account-list mm-box--display-flex mm-box--flex-direction-column mm-box--width-full"
-              data-testid="custody-account-list"
-            >
-              <div
-                class="mm-box custody-account-list__item mm-box--display-flex"
-              >
-                <div
-                  class="mm-box mm-box--display-flex mm-box--align-items-flex-start"
-                  data-testid="custody-account-list-item-radio-button"
-                />
-                <div
-                  class="mm-box mm-box--margin-left-2 mm-box--display-flex mm-box--flex-direction-column mm-box--width-full"
-                >
-                  <label
-                    class="box mm-text mm-label mm-label--html-for custody-account-list__item__title mm-text--body-md mm-text--font-weight-medium box--margin-top-2 box--margin-left-2 box--display-flex box--flex-direction-row box--align-items-center box--color-text-default"
-                    for="address-0"
-                  >
-                    <span
-                      class="box mm-text custody-account-list__item__name mm-text--inherit box--padding-right-1 box--flex-direction-row box--color-text-default"
-                    >
-                      Test name account
-                    </span>
-                  </label>
-                  <label
-                    class="box mm-text mm-label mm-label--html-for mm-text--body-md mm-text--font-weight-medium box--margin-top-2 box--margin-right-3 box--margin-left-2 box--display-flex box--flex-direction-row box--align-items-center box--color-text-default"
-                    for="address-0"
-                  >
-                    <span
-                      class="box mm-text custody-account-list__item mm-text--body-md box--display-flex box--flex-direction-row box--color-text-default"
-                    >
-                      <a
-                        class="box mm-text mm-button-base mm-button-link mm-button-link--size-auto mm-text--body-md-medium box--display-inline-flex box--flex-direction-row box--justify-content-center box--align-items-center box--color-primary-default box--background-color-transparent"
-                        href="https://etherscan.io/address/0xaD6D458402F60fD3Bd25163575031ACDce07538D"
-                        rel="noopener noreferrer"
-                        target="_blank"
-                      >
-                        <span
-                          class="box mm-text mm-text--inherit box--flex-direction-row box--color-primary-default"
-                        >
-                          0xaD6...538D
-                          <span
-                            class="box mm-icon mm-icon--size-md box--margin-left-1 box--display-inline-block box--flex-direction-row box--color-primary-default"
-                            style="mask-image: url('./images/icons/undefined.svg');"
-                          />
-                        </span>
-                      </a>
-                      <div>
-                        <div
-                          aria-describedby="tippy-tooltip-1"
-                          class=""
-                          data-original-title="Copy to clipboard"
-                          data-tooltipped=""
-                          style="display: inline; background-color: transparent;"
-                          tabindex="0"
-                        >
-                          <button
-                            class="custody-account-list__item__clipboard"
-                          >
-                            <span
-                              class="box mm-icon mm-icon--size-md box--display-inline-block box--flex-direction-row box--color-icon-muted"
-                              style="mask-image: url('./images/icons/undefined.svg');"
-                            />
-                          </button>
-                        </div>
-                      </div>
-                    </span>
-                  </label>
-                  <div
-                    class="mm-box mm-box--display-flex mm-box--justify-content-space-between"
-                  >
-                    <label
-                      class="box mm-text mm-label mm-label--html-for mm-text--body-md mm-text--font-weight-medium box--display-flex box--flex-direction-row box--align-items-center box--color-text-default"
-                      for="address-0"
-                    />
-                  </div>
-                </div>
-              </div>
-            </div>
-          </div>
-        </div>
-      </div>
-    </div>
-    <div
-      class="modal-container__footer"
-    >
-      <button
-        class="button btn--rounded btn-secondary modal-container__footer-button"
-        role="button"
-        tabindex="0"
-      >
-        Nevermind
-      </button>
-      <button
-        class="button btn--rounded btn-primary modal-container__footer-button"
-        role="button"
-        tabindex="0"
-      >
-        Remove
-      </button>
-    </div>
-  </div>
-</div>
-`;
+exports[`Confirm Remove JWT should render correctly 1`] = `<div />`;
diff --git a/ui/components/institutional/confirm-remove-jwt-modal/confirm-remove-jwt-modal.js b/ui/components/institutional/confirm-remove-jwt-modal/confirm-remove-jwt-modal.js
index 1f82c384a..aa0439b39 100644
--- a/ui/components/institutional/confirm-remove-jwt-modal/confirm-remove-jwt-modal.js
+++ b/ui/components/institutional/confirm-remove-jwt-modal/confirm-remove-jwt-modal.js
@@ -1,16 +1,24 @@
 import React, { memo, useState, useEffect } from 'react';
 import PropTypes from 'prop-types';
 import { useDispatch } from 'react-redux';
-import Modal from '../../app/modal';
 import CustodyAccountList from '../../../pages/institutional/connect-custody/account-list';
 import { useI18nContext } from '../../../hooks/useI18nContext';
 import { removeAccount } from '../../../store/actions';
 import withModalProps from '../../../helpers/higher-order-components/with-modal-props';
-import { Text } from '../../component-library/text/deprecated';
-import Box from '../../ui/box';
+import {
+  Box,
+  Text,
+  Modal,
+  ModalContent,
+  ModalHeader,
+  ModalOverlay,
+  Button,
+  BUTTON_VARIANT,
+  BUTTON_SIZES,
+} from '../../component-library';
 import {
   BorderRadius,
-  DISPLAY,
+  Display,
   TextAlign,
   TextColor,
   TextVariant,
@@ -19,16 +27,25 @@ import {
 const ConfirmRemoveJWT = ({
   custodyAccountDetails,
   accounts,
-  token,
+  token: propsToken,
   hideModal,
 }) => {
   const t = useI18nContext();
   const dispatch = useDispatch();
   const [showMore, setShowMore] = useState(false);
   const [tokenAccounts, setTokenAccounts] = useState([]);
+  let token = null;
+
+  if (propsToken) {
+    if (typeof propsToken === 'object') {
+      token = propsToken.address;
+    } else {
+      token = propsToken;
+    }
+  }
 
   useEffect(() => {
-    const lowercasedTokenAddress = token.address.toLowerCase();
+    const lowercasedTokenAddress = token.toLowerCase();
 
     const filteredAccounts = custodyAccountDetails.filter((item) => {
       const addressLower = item.address.toLowerCase();
@@ -72,52 +89,60 @@ const ConfirmRemoveJWT = ({
   };
 
   return (
-    <Modal
-      headerText={`${t('removeJWT')}?`}
-      onClose={hideModal}
-      onSubmit={handleRemove}
-      onCancel={hideModal}
-      submitText={t('remove')}
-      cancelText={t('nevermind')}
-      submitType="primary"
-    >
-      <Box
-        display={DISPLAY.FLEX}
-        padding={2}
-        borderRadius={BorderRadius.SM}
-        className="confirm-action-jwt__jwt"
-      >
-        {showMore && token ? token.address : `...${token.address.slice(-9)}`}
-      </Box>
-      {!showMore && (
-        <Text
-          color={TextColor.goerli}
-          marginLeft={2}
-          className="confirm-action-jwt__show-more"
+    <Modal isOpen onClose={hideModal}>
+      <ModalOverlay />
+      <ModalContent>
+        <ModalHeader onClose={hideModal}>{t('removeJWT')}</ModalHeader>
+        <Box
+          display={Display.Flex}
+          padding={2}
+          borderRadius={BorderRadius.SM}
+          className="confirm-action-jwt__jwt"
         >
-          <a rel="noopener noreferrer" onClick={handleShowMore}>
-            {t('showMore')}
-          </a>
+          <Text ellipsis>
+            {showMore && token ? token : `...${token.slice(-9)}`}
+          </Text>
+        </Box>
+        {!showMore && (
+          <Text
+            color={TextColor.goerli}
+            marginLeft={2}
+            className="confirm-action-jwt__show-more"
+          >
+            <a rel="noopener noreferrer" onClick={handleShowMore}>
+              {t('showMore')}
+            </a>
+          </Text>
+        )}
+        <Text
+          as="h6"
+          textAlign={TextAlign.Center}
+          variant={TextVariant.bodySm}
+          marginTop={2}
+        >
+          {t('removeJWTDescription')}
         </Text>
-      )}
-      <Text
-        as="h6"
-        textAlign={TextAlign.Center}
-        variant={TextVariant.bodySm}
-        marginTop={2}
-      >
-        {t('removeJWTDescription')}
-      </Text>
-      <Box className="confirm-action-jwt__accounts-list">
-        <CustodyAccountList accounts={tokenAccounts} rawList />
-      </Box>
+        <Box className="confirm-action-jwt__accounts-list">
+          <CustodyAccountList accounts={tokenAccounts} rawList />
+        </Box>
+        <Box display={Display.Flex}>
+          <Button
+            block
+            variant={BUTTON_VARIANT.PRIMARY}
+            size={BUTTON_SIZES.LG}
+            onClick={handleRemove}
+          >
+            {t('remove')}
+          </Button>
+        </Box>
+      </ModalContent>
     </Modal>
   );
 };
 
 ConfirmRemoveJWT.propTypes = {
   hideModal: PropTypes.func.isRequired,
-  token: PropTypes.object.isRequired,
+  token: PropTypes.oneOfType([PropTypes.string, PropTypes.object]).isRequired,
   custodyAccountDetails: PropTypes.array.isRequired,
   accounts: PropTypes.array.isRequired,
 };
diff --git a/ui/components/institutional/confirm-remove-jwt-modal/confirm-remove-jwt-modal.stories.js b/ui/components/institutional/confirm-remove-jwt-modal/confirm-remove-jwt-modal.stories.js
index ae5590c14..440541fcb 100644
--- a/ui/components/institutional/confirm-remove-jwt-modal/confirm-remove-jwt-modal.stories.js
+++ b/ui/components/institutional/confirm-remove-jwt-modal/confirm-remove-jwt-modal.stories.js
@@ -11,7 +11,7 @@ export default {
     removeAccount: () => {
       /**/
     },
-    token: { address: '0xaD6D458402F60fD3Bd25163575031ACDce07538D' },
+    token: '0xaD6D458402F60fD3Bd25163575031ACDce07538D',
     custodyAccountDetails: [
       {
         address: '0xaD6D458402F60fD3Bd25163575031ACDce07538D',
diff --git a/ui/components/institutional/confirm-remove-jwt-modal/confirm-remove-jwt-modal.test.js b/ui/components/institutional/confirm-remove-jwt-modal/confirm-remove-jwt-modal.test.js
index 4416ede93..822303f19 100644
--- a/ui/components/institutional/confirm-remove-jwt-modal/confirm-remove-jwt-modal.test.js
+++ b/ui/components/institutional/confirm-remove-jwt-modal/confirm-remove-jwt-modal.test.js
@@ -3,7 +3,7 @@ import { fireEvent, waitFor } from '@testing-library/react';
 import configureMockStore from 'redux-mock-store';
 import thunk from 'redux-thunk';
 import { renderWithProvider } from '../../../../test/lib/render-helpers';
-import testData from '../../../../.storybook/test-data';
+import mockState from '../../../../test/data/mock-state.json';
 import ConfirmRemoveJwt from '.';
 
 const mockedRemoveAccount = jest.fn();
@@ -18,7 +18,7 @@ const address = '0xaD6D458402F60fD3Bd25163575031ACDce07538D';
 
 const props = {
   hideModal: mockedHideModal,
-  token: { address },
+  token: address,
   custodyAccountDetails: [
     {
       address,
@@ -30,13 +30,8 @@ const props = {
   accounts: [{ address, balance: '0x0' }],
 };
 
-const mockStore = {
-  ...testData,
-  metamask: {},
-};
-
 const middleware = [thunk];
-const store = configureMockStore(middleware)(mockStore);
+const store = configureMockStore(middleware)(mockState);
 
 const render = () => {
   return renderWithProvider(<ConfirmRemoveJwt {...props} />, store);
diff --git a/ui/components/institutional/confirm-remove-jwt-modal/index.scss b/ui/components/institutional/confirm-remove-jwt-modal/index.scss
index 1aff21359..666ec5cad 100644
--- a/ui/components/institutional/confirm-remove-jwt-modal/index.scss
+++ b/ui/components/institutional/confirm-remove-jwt-modal/index.scss
@@ -9,5 +9,6 @@
 
   &__accounts-list {
     margin-left: -5%;
+    overflow: hidden;
   }
 }
diff --git a/ui/components/institutional/interactive-replacement-token-notification/interactive-replacement-token-notification.js b/ui/components/institutional/interactive-replacement-token-notification/interactive-replacement-token-notification.js
index 4a17f7c77..babbecd1a 100644
--- a/ui/components/institutional/interactive-replacement-token-notification/interactive-replacement-token-notification.js
+++ b/ui/components/institutional/interactive-replacement-token-notification/interactive-replacement-token-notification.js
@@ -53,7 +53,7 @@ const InteractiveReplacementTokenNotification = ({ isVisible }) => {
         return;
       }
 
-      const token = await dispatch(mmiActions.getCustodianToken());
+      const token = await dispatch(mmiActions.getCustodianToken(address));
       const custodyAccountDetails = await dispatch(
         mmiActions.getAllCustodianAccountsWithToken(
           keyring.type.split(' - ')[1],
diff --git a/ui/components/multichain/account-list-item-menu/account-list-item-menu.js b/ui/components/multichain/account-list-item-menu/account-list-item-menu.js
index ec69a70be..68bb478fd 100644
--- a/ui/components/multichain/account-list-item-menu/account-list-item-menu.js
+++ b/ui/components/multichain/account-list-item-menu/account-list-item-menu.js
@@ -233,7 +233,9 @@ export const AccountListItemMenu = ({
                 ref={removeJWTItemRef}
                 data-testid="account-options-menu__remove-jwt"
                 onClick={async () => {
-                  const token = await dispatch(mmiActions.getCustodianToken());
+                  const token = await dispatch(
+                    mmiActions.getCustodianToken(identity.address),
+                  );
                   const custodyAccountDetails = await dispatch(
                     mmiActions.getAllCustodianAccountsWithToken(
                       keyring.type.split(' - ')[1],
diff --git a/ui/pages/institutional/custody/custody.js b/ui/pages/institutional/custody/custody.js
index 0261eec90..d9825ead9 100644
--- a/ui/pages/institutional/custody/custody.js
+++ b/ui/pages/institutional/custody/custody.js
@@ -40,7 +40,7 @@ import {
   CUSTODY_ACCOUNT_DONE_ROUTE,
   DEFAULT_ROUTE,
 } from '../../../helpers/constants/routes';
-import { getCurrentChainId } from '../../../selectors';
+import { getCurrentChainId, getSelectedAddress } from '../../../selectors';
 import { getMMIConfiguration } from '../../../selectors/institutional/selectors';
 import CustodyAccountList from '../connect-custody/account-list';
 import JwtUrlForm from '../../../components/institutional/jwt-url-form';
@@ -72,6 +72,7 @@ const CustodyPage = () => {
   const [chainId, setChainId] = useState(0);
   const [connectRequest, setConnectRequest] = useState(undefined);
   const [accounts, setAccounts] = useState();
+  const address = useSelector(getSelectedAddress);
 
   const custodianButtons = useMemo(() => {
     const custodianItems = [];
@@ -259,7 +260,8 @@ const CustodyPage = () => {
       if (Object.keys(connectRequestValue).length) {
         setConnectRequest(connectRequestValue);
         setCurrentJwt(
-          connectRequestValue.token || dispatch(mmiActions.getCustodianToken()),
+          connectRequestValue.token ||
+            (await dispatch(mmiActions.getCustodianToken(address))),
         );
         setSelectedCustodianType(connectRequestValue.custodianType);
         setSelectedCustodianName(connectRequestValue.custodianName);