Wholesale replace the tokenByAddress and de-dup properly

1 files changed, 8 insertions, 4 deletions

diff --git a/packages/website/ts/local_storage/tracked_token_storage.ts b/packages/website/ts/local_storage/tracked_token_storage.ts
index 7733e8436..f92dd6298 100644
--- a/packages/website/ts/local_storage/tracked_token_storage.ts
+++ b/packages/website/ts/local_storage/tracked_token_storage.ts
@@ -1,6 +1,6 @@
 import * as _ from 'lodash';
 import { localStorage } from 'ts/local_storage/local_storage';
-import { Token, TrackedTokensByUserAddress } from 'ts/types';
+import { Token, TokenByAddress, TrackedTokensByUserAddress } from 'ts/types';
 import { configs } from 'ts/utils/configs';
 
 const TRACKED_TOKENS_KEY = 'trackedTokens';
@@ -39,10 +39,11 @@ export const trackedTokenStorage = {
         const trackedTokensByUserAddress = JSON.parse(trackedTokensJSONString);
         return trackedTokensByUserAddress;
     },
-    getTrackedTokensIfExists(userAddress: string, networkId: number): Token[] {
+    getTrackedTokensByAddress(userAddress: string, networkId: number): TokenByAddress {
+        const trackedTokensByAddress: TokenByAddress = {};
         const trackedTokensJSONString = localStorage.getItemIfExists(TRACKED_TOKENS_KEY);
         if (_.isEmpty(trackedTokensJSONString)) {
-            return undefined;
+            return trackedTokensByAddress;
         }
         const trackedTokensByUserAddress = JSON.parse(trackedTokensJSONString);
         const trackedTokensByNetworkId = trackedTokensByUserAddress[userAddress];
@@ -50,7 +51,10 @@ export const trackedTokenStorage = {
             return undefined;
         }
         const trackedTokens = trackedTokensByNetworkId[networkId];
-        return trackedTokens;
+        _.each(trackedTokens, (trackedToken: Token) => {
+            trackedTokensByAddress[trackedToken.address] = trackedToken;
+        });
+        return trackedTokensByAddress;
     },
     removeTrackedToken(userAddress: string, networkId: number, tokenAddress: string): void {
         const trackedTokensByUserAddress = this.getTrackedTokensByUserAddress();