support new socket events for token balances (#1082)

base implementation

lib/socket/types.ts

 import type { Channel } from 'phoenix';
 
-import type { AddressCoinBalanceHistoryItem } from 'types/api/address';
+import type { AddressCoinBalanceHistoryItem, AddressTokensBalancesSocketMessage } from 'types/api/address';
 import type { NewBlockSocketResponse } from 'types/api/block';
 import type { SmartContractVerificationResponse } from 'types/api/contract';
 import type { RawTracesResponse } from 'types/api/rawTrace';
@@ -18,6 +18,9 @@ SocketMessage.NewDeposits |
 SocketMessage.AddressBalance |
 SocketMessage.AddressCurrentCoinBalance |
 SocketMessage.AddressTokenBalance |
+SocketMessage.AddressTokenBalancesErc20 |
+SocketMessage.AddressTokenBalancesErc721 |
+SocketMessage.AddressTokenBalancesErc1155 |
 SocketMessage.AddressCoinBalance |
 SocketMessage.AddressTxs |
 SocketMessage.AddressTxsPending |
@@ -49,6 +52,9 @@ export namespace SocketMessage {
   export type AddressCurrentCoinBalance =
   SocketMessageParamsGeneric<'current_coin_balance', { coin_balance: string; block_number: number; exchange_rate: string }>;
   export type AddressTokenBalance = SocketMessageParamsGeneric<'token_balance', { block_number: number }>;
+  export type AddressTokenBalancesErc20 = SocketMessageParamsGeneric<'updated_token_balances_erc_20', AddressTokensBalancesSocketMessage>;
+  export type AddressTokenBalancesErc721 = SocketMessageParamsGeneric<'updated_token_balances_erc_721', AddressTokensBalancesSocketMessage>;
+  export type AddressTokenBalancesErc1155 = SocketMessageParamsGeneric<'updated_token_balances_erc_1155', AddressTokensBalancesSocketMessage>;
   export type AddressCoinBalance = SocketMessageParamsGeneric<'coin_balance', { coin_balance: AddressCoinBalanceHistoryItem }>;
   export type AddressTxs = SocketMessageParamsGeneric<'transaction', { transactions: Array<Transaction> }>;
   export type AddressTxsPending = SocketMessageParamsGeneric<'pending_transaction', { transactions: Array<Transaction> }>;

playwright/fixtures/socketServer.ts

@@ -2,7 +2,7 @@ import type { TestFixture, Page } from '@playwright/test';
 import type { WebSocket } from 'ws';
 import { WebSocketServer } from 'ws';
 
-import type { AddressCoinBalanceHistoryItem } from 'types/api/address';
+import type { AddressCoinBalanceHistoryItem, AddressTokensBalancesSocketMessage } from 'types/api/address';
 import type { NewBlockSocketResponse } from 'types/api/block';
 import type { SmartContractVerificationResponse } from 'types/api/contract';
 import type { TokenTransfer } from 'types/api/tokenTransfer';
@@ -59,6 +59,9 @@ export const joinChannel = async(socket: WebSocket, channelName: string) => {
 
 export function sendMessage(socket: WebSocket, channel: Channel, msg: 'coin_balance', payload: { coin_balance: AddressCoinBalanceHistoryItem }): void;
 export function sendMessage(socket: WebSocket, channel: Channel, msg: 'token_balance', payload: { block_number: number }): void;
+export function sendMessage(socket: WebSocket, channel: Channel, msg: 'updated_token_balances_erc_20', payload: AddressTokensBalancesSocketMessage): void;
+export function sendMessage(socket: WebSocket, channel: Channel, msg: 'updated_token_balances_erc_721', payload: AddressTokensBalancesSocketMessage): void;
+export function sendMessage(socket: WebSocket, channel: Channel, msg: 'updated_token_balances_erc_1155', payload: AddressTokensBalancesSocketMessage): void;
 export function sendMessage(socket: WebSocket, channel: Channel, msg: 'transaction', payload: { transaction: number }): void;
 export function sendMessage(socket: WebSocket, channel: Channel, msg: 'transaction', payload: { transactions: Array<Transaction> }): void;
 export function sendMessage(socket: WebSocket, channel: Channel, msg: 'pending_transaction', payload: { pending_transaction: number }): void;

playwright/index.ts

@@ -8,6 +8,7 @@ import * as router from 'next/router';
 const NEXT_ROUTER_MOCK = {
   query: {},
   pathname: '',
+  push: () => Promise.resolve(),
 };
 
 beforeMount(async({ hooksConfig }) => {

types/api/address.ts

@@ -59,6 +59,11 @@ export interface AddressTokensResponse {
   } | null;
 }
 
+export interface AddressTokensBalancesSocketMessage {
+  overflow: boolean;
+  token_balances: Array<AddressTokenBalance>;
+}
+
 export interface AddressTransactionsResponse {
   items: Array<Transaction>;
   next_page_params: {

ui/address/AddressTokens.pw.tsx

@@ -2,14 +2,15 @@ import { Box } from '@chakra-ui/react';
 import { test as base, expect, devices } from '@playwright/experimental-ct-react';
 import React from 'react';
 
-import { withName } from 'mocks/address/address';
+import * as addressMock from 'mocks/address/address';
 import * as tokensMock from 'mocks/address/tokens';
+import * as socketServer from 'playwright/fixtures/socketServer';
 import TestApp from 'playwright/TestApp';
 import buildApiUrl from 'playwright/utils/buildApiUrl';
 
 import AddressTokens from './AddressTokens';
 
-const ADDRESS_HASH = withName.hash;
+const ADDRESS_HASH = addressMock.withName.hash;
 const API_URL_ADDRESS = buildApiUrl('address', { hash: ADDRESS_HASH });
 const API_URL_TOKENS = buildApiUrl('address_tokens', { hash: ADDRESS_HASH });
 
@@ -37,7 +38,7 @@ const test = base.extend({
 
     await page.route(API_URL_ADDRESS, (route) => route.fulfill({
       status: 200,
-      body: JSON.stringify(withName),
+      body: JSON.stringify(addressMock.withName),
     }));
     await page.route(API_URL_TOKENS + '?type=ERC-20', (route) => route.fulfill({
       status: 200,
@@ -173,3 +174,106 @@ test.describe('mobile', () => {
     await expect(component).toHaveScreenshot();
   });
 });
+
+base.describe('update balances via socket', () => {
+  const test = base.extend<socketServer.SocketServerFixture>({
+    createSocket: socketServer.createSocket,
+  });
+  test.describe.configure({ mode: 'serial' });
+
+  test('', async({ mount, page, createSocket }) => {
+    test.slow();
+
+    const hooksConfig = {
+      router: {
+        query: { hash: ADDRESS_HASH, tab: 'tokens_erc20' },
+        isReady: true,
+      },
+    };
+
+    const response20 = {
+      items: [ tokensMock.erc20a, tokensMock.erc20b ],
+      next_page_params: null,
+    };
+    const response721 = {
+      items: [ tokensMock.erc721a, tokensMock.erc721b ],
+      next_page_params: null,
+    };
+    const response1155 = {
+      items: [ tokensMock.erc1155a ],
+      next_page_params: null,
+    };
+
+    await page.route(API_URL_ADDRESS, (route) => route.fulfill({
+      status: 200,
+      body: JSON.stringify(addressMock.validator),
+    }));
+    await page.route(API_URL_TOKENS + '?type=ERC-20', (route) => route.fulfill({
+      status: 200,
+      body: JSON.stringify(response20),
+    }));
+    await page.route(API_URL_TOKENS + '?type=ERC-721', (route) => route.fulfill({
+      status: 200,
+      body: JSON.stringify(response721),
+    }));
+    await page.route(API_URL_TOKENS + '?type=ERC-1155', (route) => route.fulfill({
+      status: 200,
+      body: JSON.stringify(response1155),
+    }));
+
+    const component = await mount(
+      <TestApp withSocket>
+        <Box>
+          <Box h={{ base: '134px', lg: 6 }}/>
+          <AddressTokens/>
+        </Box>
+      </TestApp>,
+      { hooksConfig },
+    );
+
+    await page.waitForResponse(API_URL_TOKENS + '?type=ERC-20');
+    await page.waitForResponse(API_URL_TOKENS + '?type=ERC-721');
+    await page.waitForResponse(API_URL_TOKENS + '?type=ERC-1155');
+
+    await expect(component).toHaveScreenshot();
+
+    const socket = await createSocket();
+    const channel = await socketServer.joinChannel(socket, `addresses:${ ADDRESS_HASH.toLowerCase() }`);
+    socketServer.sendMessage(socket, channel, 'updated_token_balances_erc_20', {
+      overflow: false,
+      token_balances: [
+        {
+          ...tokensMock.erc20a,
+          token: {
+            ...tokensMock.erc20a.token,
+            exchange_rate: '0.01',
+          },
+        },
+        {
+          ...tokensMock.erc20c,
+          value: '9852000000000000',
+          token: {
+            ...tokensMock.erc20c.token,
+            address: '0xE2cf36D00C57e01371b94B4206ae2CF841931Adc',
+            name: 'Tether USD',
+            symbol: 'USDT',
+          },
+        },
+      ],
+    });
+    socketServer.sendMessage(socket, channel, 'updated_token_balances_erc_721', {
+      overflow: false,
+      token_balances: [
+        {
+          ...tokensMock.erc721c,
+          token: {
+            ...tokensMock.erc721c.token,
+            exchange_rate: '20',
+          },
+        },
+      ],
+    });
+
+    await expect(component).toHaveScreenshot();
+  });
+});

ui/address/AddressTokens.tsx

 import { Box } from '@chakra-ui/react';
+import { useQueryClient } from '@tanstack/react-query';
 import { useRouter } from 'next/router';
 import React from 'react';
 
+import type { SocketMessage } from 'lib/socket/types';
+import type { AddressTokenBalance, AddressTokensBalancesSocketMessage, AddressTokensResponse } from 'types/api/address';
 import type { TokenType } from 'types/api/token';
 import type { PaginationParams } from 'ui/shared/pagination/types';
 
+import { getResourceKey } from 'lib/api/useApiQuery';
 import useIsMobile from 'lib/hooks/useIsMobile';
+import getQueryParamString from 'lib/router/getQueryParamString';
+import useSocketChannel from 'lib/socket/useSocketChannel';
+import useSocketMessage from 'lib/socket/useSocketMessage';
 import { ADDRESS_TOKEN_BALANCE_ERC_1155, ADDRESS_TOKEN_BALANCE_ERC_20, ADDRESS_TOKEN_BALANCE_ERC_721 } from 'stubs/address';
 import { generateListStub } from 'stubs/utils';
 import { tokenTabsByType } from 'ui/pages/Address';
@@ -30,51 +37,116 @@ const TAB_LIST_PROPS_MOBILE = {
   columnGap: 3,
 };
 
+const tokenBalanceItemIdentityFactory = (match: AddressTokenBalance) => (item: AddressTokenBalance) => ((
+  match.token.address === item.token.address &&
+  match.token_id === item.token_id &&
+  match.token_instance?.id === item.token_instance?.id
+));
+
 const AddressTokens = () => {
   const router = useRouter();
   const isMobile = useIsMobile();
 
   const scrollRef = React.useRef<HTMLDivElement>(null);
 
-  const tab = router.query.tab?.toString();
-  const tokenType: TokenType = (Object.keys(tokenTabsByType) as Array<TokenType>).find(key => tokenTabsByType[key] === tab) || 'ERC-20';
+  const tab = getQueryParamString(router.query.tab);
+  const hash = getQueryParamString(router.query.hash);
 
   const erc20Query = useQueryWithPages({
     resourceName: 'address_tokens',
-    pathParams: { hash: router.query.hash?.toString() },
+    pathParams: { hash },
     filters: { type: 'ERC-20' },
     scrollRef,
     options: {
       refetchOnMount: false,
-      enabled: tokenType === 'ERC-20',
       placeholderData: generateListStub<'address_tokens'>(ADDRESS_TOKEN_BALANCE_ERC_20, 10, { next_page_params: null }),
     },
   });
 
   const erc721Query = useQueryWithPages({
     resourceName: 'address_tokens',
-    pathParams: { hash: router.query.hash?.toString() },
+    pathParams: { hash },
     filters: { type: 'ERC-721' },
     scrollRef,
     options: {
       refetchOnMount: false,
-      enabled: tokenType === 'ERC-721',
       placeholderData: generateListStub<'address_tokens'>(ADDRESS_TOKEN_BALANCE_ERC_721, 10, { next_page_params: null }),
     },
   });
 
   const erc1155Query = useQueryWithPages({
     resourceName: 'address_tokens',
-    pathParams: { hash: router.query.hash?.toString() },
+    pathParams: { hash },
     filters: { type: 'ERC-1155' },
     scrollRef,
     options: {
       refetchOnMount: false,
-      enabled: tokenType === 'ERC-1155',
       placeholderData: generateListStub<'address_tokens'>(ADDRESS_TOKEN_BALANCE_ERC_1155, 10, { next_page_params: null }),
     },
   });
 
+  const queryClient = useQueryClient();
+
+  const updateTokensData = React.useCallback((type: TokenType, payload: AddressTokensBalancesSocketMessage) => {
+    const queryKey = getResourceKey('address_tokens', { pathParams: { hash }, queryParams: { type } });
+
+    queryClient.setQueryData(queryKey, (prevData: AddressTokensResponse | undefined) => {
+      const items = prevData?.items.map((currentItem) => {
+        const updatedData = payload.token_balances.find(tokenBalanceItemIdentityFactory(currentItem));
+        return updatedData ?? currentItem;
+      }) || [];
+
+      const extraItems = prevData?.next_page_params ?
+        [] :
+        payload.token_balances.filter((socketItem) => !items.some(tokenBalanceItemIdentityFactory(socketItem)));
+
+      if (!prevData) {
+        return {
+          items: extraItems,
+          next_page_params: null,
+        };
+      }
+
+      return {
+        items: items.concat(extraItems),
+        next_page_params: prevData.next_page_params,
+      };
+    });
+  }, [ hash, queryClient ]);
+
+  const handleTokenBalancesErc20Message: SocketMessage.AddressTokenBalancesErc20['handler'] = React.useCallback((payload) => {
+    updateTokensData('ERC-20', payload);
+  }, [ updateTokensData ]);
+
+  const handleTokenBalancesErc721Message: SocketMessage.AddressTokenBalancesErc721['handler'] = React.useCallback((payload) => {
+    updateTokensData('ERC-721', payload);
+  }, [ updateTokensData ]);
+
+  const handleTokenBalancesErc1155Message: SocketMessage.AddressTokenBalancesErc1155['handler'] = React.useCallback((payload) => {
+    updateTokensData('ERC-1155', payload);
+  }, [ updateTokensData ]);
+
+  const channel = useSocketChannel({
+    topic: `addresses:${ hash.toLowerCase() }`,
+    isDisabled: erc20Query.isPlaceholderData || erc721Query.isPlaceholderData || erc1155Query.isPlaceholderData,
+  });
+
+  useSocketMessage({
+    channel,
+    event: 'updated_token_balances_erc_20',
+    handler: handleTokenBalancesErc20Message,
+  });
+  useSocketMessage({
+    channel,
+    event: 'updated_token_balances_erc_721',
+    handler: handleTokenBalancesErc721Message,
+  });
+  useSocketMessage({
+    channel,
+    event: 'updated_token_balances_erc_1155',
+    handler: handleTokenBalancesErc1155Message,
+  });
+
   const tabs = [
     { id: tokenTabsByType['ERC-20'], title: 'ERC-20', component: <ERC20Tokens tokensQuery={ erc20Query }/> },
     { id: tokenTabsByType['ERC-721'], title: 'ERC-721', component: <ERC721Tokens tokensQuery={ erc721Query }/> },