Bugfix/update custom token balance

extension/e2e-tests/addAsset.test.ts

@@ -1,5 +1,5 @@
 import { test, expect, expectPageToHaveScreenshot } from "./test-fixtures";
-import { loginToTestAccount } from "./helpers/login";
+import { loginAndFund, loginToTestAccount } from "./helpers/login";
 import { TEST_TOKEN_ADDRESS, USDC_TOKEN_ADDRESS } from "./helpers/test-token";
 import {
   stubAccountBalances,
@@ -48,6 +48,39 @@ test("Adding unverified Soroban token", async ({ page, extensionId }) => {
       "Expected token to be either on or not on lists, but neither was visible",
     );
   }
+  await page.getByTestId("ManageAssetRowButton").click();
+  await expect(page.getByTestId("ToggleToken__asset-code")).toHaveText(
+    "E2E Token",
+  );
+  await expect(page.getByTestId("ToggleToken__asset-add-remove")).toHaveText(
+    "Add Token",
+  );
+  await page.getByRole("button", { name: "Confirm" }).click();
+  await expect(
+    page.getByTestId("ManageAssetRowButton__ellipsis-E2E"),
+  ).toBeVisible();
+
+  // now go back and remove this asset
+  await page.getByTestId("BackButton").click();
+  await expect(page.getByText("Your assets")).toBeVisible();
+  await expect(page.getByTestId("ManageAssetCode")).toHaveText("E2E");
+  await expect(page.getByTestId("ManageAssetDomain")).toHaveText(
+    "Stellar Network",
+  );
+  await page.getByTestId("ManageAssetRowButton__ellipsis-E2E").click();
+  await page.getByText("Remove asset").click();
+  await expect(page.getByTestId("ToggleToken__asset-code")).toHaveText(
+    "CBVX…HWXJ",
+  );
+  await expect(page.getByTestId("ToggleToken__asset-add-remove")).toHaveText(
+    "Remove Token",
+  );
+  await page.getByRole("button", { name: "Confirm" }).click();
+  await expect(
+    page.getByText(
+      "You have no assets added. Get started by adding an asset below.",
+    ),
+  ).toBeVisible();
 });
 
 // Skipping this test because on Testnet, stellar.expert's asset list is formatter incorrectly
@@ -116,6 +149,51 @@ test("Adding token on Futurenet", async ({ page, extensionId, context }) => {
   await page.getByText("Add an asset").click({ force: true });
   await expect(page.getByTestId("search-token-input")).toBeVisible();
 });
+test("Adding classic asset on Testnet", async ({ page, extensionId }) => {
+  test.slow();
+  await loginAndFund({ page, extensionId });
+
+  await page.getByTestId("account-options-dropdown").click();
+  await page.getByText("Manage assets").click();
+  await expect(page.getByText("Your assets")).toBeVisible();
+  await page.getByText("Add an asset").click({ force: true });
+  await page
+    .getByTestId("search-asset-input")
+    .fill("GBBD47IF6LWK7P7MDEVSCWR7DPUWV3NY3DTQEVFL4NAT4AQH3ZLLFLA5");
+  await expect(page.getByText("USDC")).toBeVisible();
+
+  await page.getByTestId("ManageAssetRowButton").click();
+  await expect(
+    page.getByTestId("SignTransaction__TrustlineRow__Asset"),
+  ).toHaveText("USDC");
+  await expect(
+    page.getByTestId("SignTransaction__TrustlineRow__Type"),
+  ).toHaveText("Add Trustline");
+  await page.getByRole("button", { name: "Confirm" }).click();
+  await expect(
+    page.getByTestId("ManageAssetRowButton__ellipsis-USDC"),
+  ).toBeVisible();
+
+  // now go back and remove this asset
+  await page.getByTestId("BackButton").click();
+  await expect(page.getByText("Your assets")).toBeVisible();
+  await expect(page.getByTestId("ManageAssetCode")).toHaveText("USDC");
+  await expect(page.getByTestId("ManageAssetDomain")).toHaveText("centre.io");
+  await page.getByTestId("ManageAssetRowButton__ellipsis-USDC").click();
+  await page.getByText("Remove asset").click();
+  await expect(
+    page.getByTestId("SignTransaction__TrustlineRow__Asset"),
+  ).toHaveText("USDC");
+  await expect(
+    page.getByTestId("SignTransaction__TrustlineRow__Type"),
+  ).toHaveText("Remove Trustline");
+  await page.getByRole("button", { name: "Confirm" }).click();
+  await expect(
+    page.getByText(
+      "You have no assets added. Get started by adding an asset below.",
+    ),
+  ).toBeVisible();
+});
 test.afterAll(async ({ page, extensionId }) => {
   if (
     process.env.IS_INTEGRATION_MODE &&

extension/src/popup/components/manageAssets/ChooseAsset/index.tsx

@@ -6,6 +6,7 @@ import { useTranslation } from "react-i18next";
 
 import { ROUTES } from "popup/constants/routes";
 import { settingsSorobanSupportedSelector } from "popup/ducks/settings";
+import { balancesSelector } from "popup/ducks/cache";
 import { SubviewHeader } from "popup/components/SubviewHeader";
 import { View } from "popup/basics/layout/View";
 
@@ -31,6 +32,7 @@ export const ChooseAsset = ({
   const { t } = useTranslation();
   const location = useLocation();
   const isSorobanSuported = useSelector(settingsSorobanSupportedSelector);
+  const cachedBalances = useSelector(balancesSelector);
 
   const ManageAssetRowsWrapperRef = useRef<HTMLDivElement>(null);
 
@@ -44,12 +46,14 @@ export const ChooseAsset = ({
     domainState.state === RequestState.LOADING;
 
   useEffect(() => {
+    /* This effect is keyed off of changes to cachedBalances as this let's us update the UI when an asset is removed */
     const getData = async () => {
       await fetchData(true);
     };
+
     getData();
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, []);
+  }, [cachedBalances]);
 
   if (isLoading) {
     return (

extension/src/popup/components/manageAssets/ManageAssetRows/ToggleTokenInternal/index.tsx

@@ -1,20 +1,22 @@
 import React from "react";
 import { useDispatch } from "react-redux";
-import { useNavigate } from "react-router-dom";
 import { Asset, Badge, Button, Icon, Text } from "@stellar/design-system";
 import { Networks } from "stellar-sdk";
+import { captureException } from "@sentry/browser";
 
 import { KeyIdenticon } from "popup/components/identicons/KeyIdenticon";
 import { AppDispatch } from "popup/App";
 import { NetworkDetails, NETWORKS } from "@shared/constants/stellar";
 import { addTokenId } from "popup/ducks/accountServices";
 import { removeTokenId } from "popup/ducks/transactionSubmission";
-import { navigateTo } from "popup/helpers/navigate";
-import { ROUTES } from "popup/constants/routes";
 
-import "./styles.scss";
 import { isSacContract } from "popup/helpers/soroban";
-import { truncateString } from "helpers/stellar";
+import { isMainnet, truncateString } from "helpers/stellar";
+import { AccountBalances, useGetBalances } from "helpers/hooks/useGetBalances";
+
+import { isError } from "helpers/request";
+
+import "./styles.scss";
 
 interface ToggleTokenInternalProps {
   asset: {
@@ -38,7 +40,10 @@ export const ToggleTokenInternal = ({
   publicKey,
 }: ToggleTokenInternalProps) => {
   const dispatch: AppDispatch = useDispatch();
-  const nav = useNavigate();
+  const { fetchData: fetchBalances } = useGetBalances({
+    showHidden: false,
+    includeIcons: false,
+  });
 
   const onConfirm = async () => {
     if (!asset.isTrustlineActive) {
@@ -57,7 +62,23 @@ export const ToggleTokenInternal = ({
         }),
       );
     }
-    navigateTo(ROUTES.account, nav);
+    const balancesResult = await fetchBalances(
+      publicKey,
+      isMainnet(networkDetails),
+      networkDetails,
+      false,
+    );
+
+    if (isError<AccountBalances>(balancesResult)) {
+      // we don't want to throw an error if balances fail to fetch as this doesn't affect the tx submission
+      // let's simply log the error and continue - the user will need to refresh the Account page or wait for polling to refresh the balances
+      captureException(
+        `Failed to fetch balances after ${!asset.isTrustlineActive ? "add" : "remove"} token - ${JSON.stringify(
+          balancesResult.message,
+        )} ${networkDetails.network}`,
+      );
+    }
+    onCancel();
   };
   const isSac =
     !!asset.name &&
@@ -93,10 +114,18 @@ export const ToggleTokenInternal = ({
             </div>
           )}
 
-          <Text as="div" size="sm" weight="medium">
+          <Text
+            as="div"
+            size="sm"
+            weight="medium"
+            data-testid="ToggleToken__asset-code"
+          >
             {isSac ? asset.code : asset.name || truncateString(asset.contract!)}
           </Text>
-          <div className="ToggleToken__wrapper__badge">
+          <div
+            className="ToggleToken__wrapper__badge"
+            data-testid="ToggleToken__asset-add-remove"
+          >
             <Badge
               size="sm"
               variant="secondary"

extension/src/popup/components/manageAssets/SearchAsset/index.tsx

@@ -3,8 +3,8 @@ import { Navigate, useLocation } from "react-router-dom";
 import { Formik, Form, Field, FieldProps } from "formik";
 import debounce from "lodash/debounce";
 import { useTranslation } from "react-i18next";
-
 import { Notification } from "@stellar/design-system";
+import { useSelector } from "react-redux";
 
 import { FormRows } from "popup/basics/Forms";
 import { ROUTES } from "popup/constants/routes";
@@ -13,17 +13,19 @@ import { isMainnet, isTestnet } from "helpers/stellar";
 import { SubviewHeader } from "popup/components/SubviewHeader";
 import { View } from "popup/basics/layout/View";
 
-import { ManageAssetRows } from "../ManageAssetRows";
-import { SearchInput, SearchCopy, SearchResults } from "../AssetResults";
-import { useAssetLookup } from "./hooks/useAssetLookup";
-import { useGetSearchData } from "./hooks/useSearchData";
 import { NetworkDetails } from "@shared/constants/stellar";
 import { RequestState } from "helpers/hooks/fetchHookInterface";
 import { Loading } from "popup/components/Loading";
 import { openTab } from "popup/helpers/navigate";
 import { newTabHref } from "helpers/urls";
 import { AppDataType } from "helpers/hooks/useGetAppData";
 import { reRouteOnboarding } from "popup/helpers/route";
+import { balancesSelector } from "popup/ducks/cache";
+
+import { ManageAssetRows } from "../ManageAssetRows";
+import { SearchInput, SearchCopy, SearchResults } from "../AssetResults";
+import { useAssetLookup } from "./hooks/useAssetLookup";
+import { useGetSearchData } from "./hooks/useSearchData";
 
 import "./styles.scss";
 
@@ -62,6 +64,7 @@ const ResultsHeader = () => {
 export const SearchAsset = () => {
   const { t } = useTranslation();
   const location = useLocation();
+  const cachedBalances = useSelector(balancesSelector);
 
   const ResultsRef = useRef<HTMLDivElement>(null);
 
@@ -95,12 +98,13 @@ export const SearchAsset = () => {
   const handleSearch = handleSearchRef.current;
 
   useEffect(() => {
+    /* This effect is keyed off of changes to cachedBalances as this let's us update the UI when an asset is removed */
     const getData = async () => {
       await fetchData(true);
     };
 
     getData();
-  }, []);
+  }, [cachedBalances]);
 
   if (
     state.state === RequestState.IDLE ||

extension/src/popup/views/SignTransaction/index.tsx

@@ -597,10 +597,16 @@ export const Trustline = ({ operations, icons }: TrustlineProps) => {
 
     return (
       <div className="SignTransaction__TrustlineRow">
-        <div className="SignTransaction__TrustlineRow__Asset">
+        <div
+          className="SignTransaction__TrustlineRow__Asset"
+          data-testid="SignTransaction__TrustlineRow__Asset"
+        >
           {renderTrustlineAsset(line)}
         </div>
-        <div className="SignTransaction__TrustlineRow__Type">
+        <div
+          className="SignTransaction__TrustlineRow__Type"
+          data-testid="SignTransaction__TrustlineRow__Type"
+        >
           {isRemoveTrustline ? (
             <>
               <Icon.MinusCircle />