From 31bce5e97c39e648a6d78040bf8e88c489c6988c Mon Sep 17 00:00:00 2001
From: Matt Holtzman <matt.holtzman@gmail.com>
Date: Tue, 7 Feb 2023 16:31:06 -0500
Subject: [PATCH] update balances when custom token details change

---
 main/store/actions/index.js           | 20 ++++++++++++
 test/main/store/actions/index.test.js | 45 +++++++++++++++++++++++++--
 2 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/main/store/actions/index.js b/main/store/actions/index.js
index d2696dcac..08a6623f5 100644
--- a/main/store/actions/index.js
+++ b/main/store/actions/index.js
@@ -623,6 +623,26 @@ module.exports = {
 
       return [...existingTokens, ...tokensToAdd]
     })
+
+    u('main.balances', (balances) => {
+      // update the balances for any custom tokens that changed
+      Object.values(balances).forEach((accountBalances) => {
+        tokens.forEach((token) => {
+          const tokenAddress = token.address.toLowerCase()
+          const matchingBalance = accountBalances.find(
+            (b) => b.address.toLowerCase() === tokenAddress && b.chainId === token.chainId
+          )
+
+          if (matchingBalance) {
+            matchingBalance.logoURI = token.logoURI || matchingBalance.logoURI
+            matchingBalance.symbol = token.symbol || matchingBalance.symbol
+            matchingBalance.name = token.name || matchingBalance.symbol
+          }
+        })
+      })
+
+      return balances
+    })
   },
   removeCustomTokens: (u, tokens) => {
     u('main.tokens.custom', (existing) => {
diff --git a/test/main/store/actions/index.test.js b/test/main/store/actions/index.test.js
index 3cccb942f..5d9c3f0f3 100644
--- a/test/main/store/actions/index.test.js
+++ b/test/main/store/actions/index.test.js
@@ -375,12 +375,17 @@ describe('#removeBalance', () => {
 })
 
 describe('#addCustomTokens', () => {
-  let tokens = []
+  let tokens = [],
+    balances = {}
 
   const updaterFn = (node, update) => {
-    expect(node).toBe('main.tokens.custom')
+    if (node === 'main.tokens.custom') {
+      tokens = update(tokens)
+    }
 
-    tokens = update(tokens)
+    if (node === 'main.balances') {
+      balances = update(balances)
+    }
   }
 
   const addTokens = (tokensToAdd) => addCustomTokensAction(updaterFn, tokensToAdd)
@@ -407,6 +412,40 @@ describe('#addCustomTokens', () => {
     expect(tokens[0]).toEqual(testTokens.zrx)
     expect(tokens[1].symbol).toBe('BAD')
   })
+
+  it('updates an existing balance for a custom token', () => {
+    const account = '0xd0e3872f5fa8ecb49f1911f605c0da90689a484e'
+
+    balances = {
+      [account]: [
+        {
+          address: testTokens.badger.address,
+          chainId: testTokens.badger.chainId,
+          symbol: 'BDG',
+          name: 'Old Badger',
+          logoURI: 'http://logo.io'
+        }
+      ]
+    }
+
+    const updatedBadgerToken = {
+      ...testTokens.badger,
+      symbol: 'BADGER',
+      name: 'Badger Token'
+    }
+
+    addTokens([updatedBadgerToken])
+
+    expect(balances[account]).toStrictEqual([
+      {
+        address: testTokens.badger.address,
+        chainId: testTokens.badger.chainId,
+        symbol: 'BADGER',
+        name: 'Badger Token',
+        logoURI: 'http://logo.io'
+      }
+    ])
+  })
 })
 
 describe('#removeCustomTokens', () => {