chore: improvements on getting cached balances

Instead of checking all tokens when fetching balances only relevant tokens are processed.

Author: saledjenic

services/wallet/reader.go

@@ -14,6 +14,8 @@ import (
 
 	"github.com/bep/debounce"
 
+	"golang.org/x/exp/maps"
+
 	"github.com/status-im/go-wallet-sdk/pkg/tokens/types"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -354,12 +356,19 @@ func (r *Reader) GetCachedBalances(chainIDs []uint64, addresses []common.Address
 		return nil, err
 	}
 
-	allTokens, err := r.tokenManager.GetTokensByChains(chainIDs)
+	balances, err := tokensToBalancesPerChain(cachedTokens)
 	if err != nil {
 		return nil, err
 	}
 
-	balances, err := tokensToBalancesPerChain(cachedTokens)
+	tokensOfInterest := make(map[string]struct{}, 0)
+	for _, cachedToken := range cachedTokens {
+		for _, token := range cachedToken {
+			tokensOfInterest[types.TokenKey(token.TokenChainID, token.TokenAddress)] = struct{}{}
+		}
+	}
+
+	allTokens, err := r.tokenManager.GetTokensByKeys(maps.Keys(tokensOfInterest))
 	if err != nil {
 		return nil, err
 	}

services/wallet/reader_impl_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/event"
 
+	"github.com/status-im/go-wallet-sdk/pkg/tokens/types"
 	wsdktypes "github.com/status-im/go-wallet-sdk/pkg/tokens/types"
 
 	"github.com/status-im/status-go/pkg/pubsub"
@@ -410,8 +411,12 @@ func TestGetCachedBalancesInternal(t *testing.T) {
 		},
 	}
 
+	tokensOfInterest := []string{
+		types.TokenKey(cachedTokens[addresses[1]][0].TokenChainID, cachedTokens[addresses[1]][0].TokenAddress),
+	}
+
 	tokenManager.EXPECT().GetCachedBalances().Return(cachedTokens, nil)
-	tokenManager.EXPECT().GetTokensByChains(chainIDs).Return(allTokens, nil)
+	tokenManager.EXPECT().GetTokensByKeys(tokensOfInterest).Return(allTokens, nil)
 	tokens, err := reader.GetCachedBalances(chainIDs, addresses)
 	require.NoError(t, err)
 

services/wallet/reader_test.go

@@ -16,7 +16,6 @@ import (
 
 	"github.com/status-im/status-go/pkg/pubsub"
 	"github.com/status-im/status-go/services/wallet"
-	"github.com/status-im/status-go/services/wallet/testutils"
 	mock_token "github.com/status-im/status-go/services/wallet/token/mock/token"
 	tokentypes "github.com/status-im/status-go/services/wallet/token/types"
 	mock_tokenbalances "github.com/status-im/status-go/services/wallet/tokenbalances/mock/storage"
@@ -121,8 +120,12 @@ func TestGetCachedBalances(t *testing.T) {
 		},
 	}
 
+	tokensOfInterest := []string{
+		types.TokenKey(cachedTokens[addresses[1]][0].TokenChainID, cachedTokens[addresses[1]][0].TokenAddress),
+	}
+
 	tokenManager.EXPECT().GetCachedBalances().Return(cachedTokens, nil)
-	tokenManager.EXPECT().GetTokensByChains(testutils.NewUint64SliceMatcher(chainIDs)).Return(allTokens, nil)
+	tokenManager.EXPECT().GetTokensByKeys(tokensOfInterest).Return(allTokens, nil)
 	tokens, err := reader.GetCachedBalances(chainIDs, addresses)
 	require.NoError(t, err)
 
@@ -191,9 +194,13 @@ func TestFetchBalances(t *testing.T) {
 		},
 	}
 
+	tokensOfInterest := []string{
+		types.TokenKey(cachedTokens[addresses[1]][0].TokenChainID, cachedTokens[addresses[1]][0].TokenAddress),
+	}
+
 	// Test GetCachedBalances with cached data
 	tokenManager.EXPECT().GetCachedBalances().Return(cachedTokens, nil)
-	tokenManager.EXPECT().GetTokensByChains(testutils.NewUint64SliceMatcher(chainIDs)).Return(allTokens, nil)
+	tokenManager.EXPECT().GetTokensByKeys(tokensOfInterest).Return(allTokens, nil)
 
 	tokens, err := reader.GetCachedBalances(chainIDs, addresses)
 	require.NoError(t, err)

services/wallet/token/token.go

@@ -68,6 +68,7 @@ type CommunityTokenImageBuilder interface {
 type ManagerInterface interface {
 	GetTokenByChainAddress(chainID uint64, address common.Address) (*tokentypes.Token, error)
 	GetTokensByChains(chainIDs []uint64) ([]*tokentypes.Token, error)
+	GetTokensByKeys(tokenKeys []string) ([]*tokentypes.Token, error)
 	GetCachedBalances() (map[common.Address][]tokentypes.StorageToken, error)
 	CacheBalances(balances map[common.Address][]tokentypes.StorageToken) error
 	FindOrCreateTokenByAddress(ctx context.Context, chainID uint64, address common.Address) (*tokentypes.Token, error)