status-go/services/wallet/token/balancefetcher/balance_fetcher_test.go

387 lines
12 KiB
Go
Raw Normal View History

package balancefetcher
import (
"context"
"errors"
"math/big"
"os"
"testing"
"github.com/golang/mock/gomock"
"github.com/stretchr/testify/require"
"github.com/ethereum/go-ethereum/accounts/abi/bind"
"github.com/ethereum/go-ethereum/common"
"github.com/ethereum/go-ethereum/common/hexutil"
"github.com/ethereum/go-ethereum/log"
"github.com/status-im/status-go/contracts/ethscan"
"github.com/status-im/status-go/params"
mock_contracts "github.com/status-im/status-go/contracts/mock"
"github.com/status-im/status-go/rpc/chain"
mock_client "github.com/status-im/status-go/rpc/chain/mock/client"
mock_network "github.com/status-im/status-go/rpc/network/mock"
w_common "github.com/status-im/status-go/services/wallet/common"
)
type FakeBalanceScanner struct {
etherBalances map[common.Address]*big.Int
tokenBalances map[common.Address]map[common.Address]*big.Int
}
func (f *FakeBalanceScanner) EtherBalances(opts *bind.CallOpts, addresses []common.Address) ([]ethscan.BalanceScannerResult, error) {
result := make([]ethscan.BalanceScannerResult, 0, len(addresses))
for _, address := range addresses {
balance, ok := f.etherBalances[address]
if !ok {
result = append(result, ethscan.BalanceScannerResult{
Success: false,
Data: []byte{},
})
} else {
result = append(result, ethscan.BalanceScannerResult{
Success: true,
Data: balance.Bytes(),
})
}
}
return result, nil
}
func (f *FakeBalanceScanner) TokenBalances(opts *bind.CallOpts, addresses []common.Address, tokenAddress common.Address) ([]ethscan.BalanceScannerResult, error) {
result := make([]ethscan.BalanceScannerResult, 0, len(addresses))
for _, address := range addresses {
balances, ok := f.tokenBalances[address]
if !ok {
result = append(result, ethscan.BalanceScannerResult{
Success: false,
Data: []byte{},
})
} else {
balance, ok := balances[tokenAddress]
if !ok {
result = append(result, ethscan.BalanceScannerResult{
Success: false,
Data: []byte{},
})
} else {
result = append(result, ethscan.BalanceScannerResult{
Success: true,
Data: balance.Bytes(),
})
}
}
}
return result, nil
}
func (f *FakeBalanceScanner) TokensBalance(opts *bind.CallOpts, owner common.Address, contracts []common.Address) ([]ethscan.BalanceScannerResult, error) {
result := make([]ethscan.BalanceScannerResult, 0, len(contracts))
for _, contract := range contracts {
balances, ok := f.tokenBalances[owner]
if !ok {
result = append(result, ethscan.BalanceScannerResult{
Success: false,
Data: []byte{},
})
} else {
balance, ok := balances[contract]
if !ok {
result = append(result, ethscan.BalanceScannerResult{
Success: false,
Data: []byte{},
})
} else {
result = append(result, ethscan.BalanceScannerResult{
Success: true,
Data: balance.Bytes(),
})
}
}
}
return result, nil
}
type FakeERC20Caller struct {
accountBalances map[common.Address]*big.Int
}
func (f *FakeERC20Caller) BalanceOf(opts *bind.CallOpts, account common.Address) (*big.Int, error) {
balance, ok := f.accountBalances[account]
if !ok {
return nil, errors.New("account not found")
}
return balance, nil
}
func (f *FakeERC20Caller) Name(opts *bind.CallOpts) (string, error) {
return "TestToken", nil
}
func (f *FakeERC20Caller) Symbol(opts *bind.CallOpts) (string, error) {
return "TT", nil
}
func (f *FakeERC20Caller) Decimals(opts *bind.CallOpts) (uint8, error) {
return 18, nil
}
func TestBalanceFetcherFetchBalancesForChainNativeAndTokensWithScanContract(t *testing.T) {
log.Root().SetHandler(log.LvlFilterHandler(log.LvlInfo, log.StreamHandler(os.Stdout, log.TerminalFormat(true))))
ctx := context.Background()
accounts := []common.Address{
common.HexToAddress("0x1234567890abcdef"),
common.HexToAddress("0xabcdef1234567890"),
}
tokens := []common.Address{
NativeChainAddress,
common.HexToAddress("0xabcdef1234567890"),
common.HexToAddress("0x0987654321fedcba"),
}
var atBlock *big.Int // nil triggers using a scan contract
ctrl := gomock.NewController(t)
defer ctrl.Finish()
networkManager := mock_network.NewMockManagerInterface(ctrl)
networkManager.EXPECT().GetAll().Return([]*params.Network{
{
ChainID: w_common.EthereumMainnet,
},
}, nil).AnyTimes()
chainClient := mock_client.NewMockClientInterface(ctrl)
chainClient.EXPECT().NetworkID().Return(w_common.EthereumMainnet).AnyTimes()
expectedEthBalances := map[common.Address]*big.Int{
accounts[0]: big.NewInt(100),
accounts[1]: big.NewInt(200),
}
expectedTokenBalances := map[common.Address]map[common.Address]*big.Int{
accounts[0]: {
tokens[1]: big.NewInt(1000),
tokens[2]: big.NewInt(2000),
},
accounts[1]: {
tokens[1]: big.NewInt(3000),
tokens[2]: big.NewInt(4000),
},
}
expectedBalances := map[common.Address]map[common.Address]*hexutil.Big{
accounts[0]: {
tokens[0]: (*hexutil.Big)(expectedEthBalances[accounts[0]]),
tokens[1]: (*hexutil.Big)(expectedTokenBalances[accounts[0]][tokens[1]]),
tokens[2]: (*hexutil.Big)(expectedTokenBalances[accounts[0]][tokens[2]]),
},
accounts[1]: {
tokens[0]: (*hexutil.Big)(expectedEthBalances[accounts[1]]),
tokens[1]: (*hexutil.Big)(expectedTokenBalances[accounts[1]][tokens[1]]),
tokens[2]: (*hexutil.Big)(expectedTokenBalances[accounts[1]][tokens[2]]),
},
}
contractMaker := mock_contracts.NewMockContractMakerIface(ctrl)
contractMaker.EXPECT().NewEthScan(w_common.EthereumMainnet).Return(&FakeBalanceScanner{
etherBalances: expectedEthBalances,
tokenBalances: expectedTokenBalances,
}, uint(0), nil).AnyTimes()
bf := NewDefaultBalanceFetcher(contractMaker)
// Fetch native balances and token balances using scan contract
balances, err := bf.fetchBalancesForChain(ctx, chainClient, accounts, tokens, atBlock)
require.NoError(t, err)
require.Equal(t, expectedBalances, balances)
}
func TestBalanceFetcherFetchBalancesForChainTokensWithTokenContracts(t *testing.T) {
log.Root().SetHandler(log.LvlFilterHandler(log.LvlInfo, log.StreamHandler(os.Stdout, log.TerminalFormat(true))))
ctx := context.Background()
accounts := []common.Address{
common.HexToAddress("0x1234567890abcdef"),
common.HexToAddress("0xabcdef1234567890"),
}
tokens := []common.Address{
common.HexToAddress("0xabcdef1234567890"),
common.HexToAddress("0x0987654321fedcba"),
}
atBlock := big.NewInt(0) // will trigger using a token contract
ctrl := gomock.NewController(t)
networkManager := mock_network.NewMockManagerInterface(ctrl)
networkManager.EXPECT().GetAll().Return([]*params.Network{
{
ChainID: w_common.EthereumMainnet,
},
}, nil).AnyTimes()
chainClient := mock_client.NewMockClientInterface(ctrl)
chainClient.EXPECT().NetworkID().Return(w_common.EthereumMainnet).AnyTimes()
chainClient.EXPECT().CallContract(gomock.Any(), gomock.Any(), atBlock).Return([]byte{}, nil).AnyTimes()
expectedTokenBalances := map[common.Address]map[common.Address]*big.Int{
tokens[0]: {
accounts[0]: big.NewInt(1000),
accounts[1]: big.NewInt(2000),
},
tokens[1]: {
accounts[0]: big.NewInt(3000),
accounts[1]: big.NewInt(4000),
},
}
expectedBalances := map[common.Address]map[common.Address]*hexutil.Big{
accounts[0]: {
tokens[0]: (*hexutil.Big)(expectedTokenBalances[tokens[0]][accounts[0]]),
tokens[1]: (*hexutil.Big)(expectedTokenBalances[tokens[1]][accounts[0]]),
},
accounts[1]: {
tokens[0]: (*hexutil.Big)(expectedTokenBalances[tokens[0]][accounts[1]]),
tokens[1]: (*hexutil.Big)(expectedTokenBalances[tokens[1]][accounts[1]]),
},
}
contractMaker := mock_contracts.NewMockContractMakerIface(ctrl)
contractMaker.EXPECT().NewEthScan(w_common.EthereumMainnet).Return(&FakeBalanceScanner{}, uint(0), nil).Times(1)
for _, token := range tokens {
contractMaker.EXPECT().NewERC20Caller(w_common.EthereumMainnet, token).Return(&FakeERC20Caller{
accountBalances: expectedTokenBalances[token],
}, nil).AnyTimes()
}
bf := NewDefaultBalanceFetcher(contractMaker)
// Fetch token balances using tokens contracts
balances, err := bf.fetchBalancesForChain(ctx, chainClient, accounts, tokens, atBlock)
require.NoError(t, err)
require.Equal(t, expectedBalances, balances)
}
func TestBalanceFetcherGetBalancesAtByChain(t *testing.T) {
log.Root().SetHandler(log.LvlFilterHandler(log.LvlInfo, log.StreamHandler(os.Stdout, log.TerminalFormat(true))))
ctx := context.Background()
accounts := []common.Address{
common.HexToAddress("0x1234567890abcdef"),
common.HexToAddress("0xabcdef1234567890"),
}
tokens := []common.Address{
NativeChainAddress,
common.HexToAddress("0xabcdef1234567890"),
common.HexToAddress("0x0987654321fedcba"),
}
var atBlock *big.Int // nil triggers using a scan contract
atBlocks := map[uint64]*big.Int{
w_common.EthereumMainnet: atBlock, // nil triggers using a scan contract
w_common.OptimismMainnet: atBlock, // nil triggers using a scan contract
}
ctrl := gomock.NewController(t)
defer ctrl.Finish()
networkManager := mock_network.NewMockManagerInterface(ctrl)
networkManager.EXPECT().GetAll().Return([]*params.Network{
{
ChainID: w_common.EthereumMainnet,
},
{
ChainID: w_common.OptimismMainnet,
},
{
ChainID: w_common.ArbitrumMainnet,
},
}, nil).AnyTimes()
chainClient := mock_client.NewMockClientInterface(ctrl)
chainClient.EXPECT().NetworkID().Return(w_common.EthereumMainnet).AnyTimes()
chainClientOpt := mock_client.NewMockClientInterface(ctrl)
chainClientOpt.EXPECT().NetworkID().Return(w_common.OptimismMainnet).AnyTimes()
chainClientArb := mock_client.NewMockClientInterface(ctrl)
chainClientArb.EXPECT().NetworkID().Return(w_common.ArbitrumMainnet).AnyTimes()
chainClients := map[uint64]chain.ClientInterface{
w_common.EthereumMainnet: chainClient,
w_common.OptimismMainnet: chainClientOpt,
w_common.ArbitrumMainnet: chainClientArb,
}
expectedEthBalances := map[common.Address]*big.Int{
accounts[0]: big.NewInt(100),
accounts[1]: big.NewInt(200),
}
expectedEthOptBalances := map[common.Address]*big.Int{
accounts[0]: big.NewInt(300),
accounts[1]: big.NewInt(400),
}
expectedTokenBalances := map[common.Address]map[common.Address]*big.Int{
accounts[0]: {
tokens[1]: big.NewInt(1000),
tokens[2]: big.NewInt(2000),
},
accounts[1]: {
tokens[1]: big.NewInt(3000),
tokens[2]: big.NewInt(4000),
},
}
expectedTokenOptBalances := map[common.Address]map[common.Address]*big.Int{
accounts[0]: {
tokens[1]: big.NewInt(5000),
tokens[2]: big.NewInt(6000),
},
}
expectedBalances := map[uint64]map[common.Address]map[common.Address]*hexutil.Big{
w_common.EthereumMainnet: {
accounts[0]: {
tokens[0]: (*hexutil.Big)(expectedEthBalances[accounts[0]]),
tokens[1]: (*hexutil.Big)(expectedTokenBalances[accounts[0]][tokens[1]]),
tokens[2]: (*hexutil.Big)(expectedTokenBalances[accounts[0]][tokens[2]]),
},
accounts[1]: {
tokens[0]: (*hexutil.Big)(expectedEthBalances[accounts[1]]),
tokens[1]: (*hexutil.Big)(expectedTokenBalances[accounts[1]][tokens[1]]),
tokens[2]: (*hexutil.Big)(expectedTokenBalances[accounts[1]][tokens[2]]),
},
},
w_common.OptimismMainnet: {
accounts[0]: {
tokens[0]: (*hexutil.Big)(expectedEthOptBalances[accounts[0]]),
tokens[1]: (*hexutil.Big)(expectedTokenOptBalances[accounts[0]][tokens[1]]),
tokens[2]: (*hexutil.Big)(expectedTokenOptBalances[accounts[0]][tokens[2]]),
},
accounts[1]: {
tokens[0]: (*hexutil.Big)(expectedEthOptBalances[accounts[1]]),
},
},
}
contractMaker := mock_contracts.NewMockContractMakerIface(ctrl)
contractMaker.EXPECT().NewEthScan(w_common.EthereumMainnet).Return(&FakeBalanceScanner{
etherBalances: expectedEthBalances,
tokenBalances: expectedTokenBalances,
}, uint(0), nil).AnyTimes()
contractMaker.EXPECT().NewEthScan(w_common.OptimismMainnet).Return(&FakeBalanceScanner{
etherBalances: expectedEthOptBalances,
tokenBalances: expectedTokenOptBalances,
}, uint(0), nil).AnyTimes()
contractMaker.EXPECT().NewEthScan(w_common.ArbitrumMainnet).Return(nil, uint(0), errors.New("no scan contract")).AnyTimes()
bf := NewDefaultBalanceFetcher(contractMaker)
// Fetch native balances and token balances using scan contract
balances, err := bf.GetBalancesAtByChain(ctx, chainClients, accounts, tokens, atBlocks)
require.NoError(t, err)
require.Equal(t, expectedBalances, balances)
}