Skip to content

Commit

Permalink
chore: add unit tests and clean up
Browse files Browse the repository at this point in the history
  • Loading branch information
@nickewansmith
nickewansmith committed Nov 1, 2024
1 parent a459bdd commit cb9a55c
Show file tree
Hide file tree
Showing 4 changed files with 544 additions and 7 deletions.
325 changes: 324 additions & 1 deletion packages/assets-controllers/src/AccountTrackerController.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -248,6 +248,123 @@ describe('AccountTrackerController', () => {
},
);
});

it('should update staked balance when includeStakedAssets is enabled', async () => {
mockedQuery
.mockReturnValueOnce(Promise.resolve('0x10'))
.mockReturnValueOnce(Promise.resolve('0x11'));

await withController(
{
options: {
includeStakedAssets: true,
getStakedBalanceForChain: jest.fn().mockResolvedValue('0x1'),
},
isMultiAccountBalancesEnabled: false,
selectedAccount: ACCOUNT_1,
listAccounts: [ACCOUNT_1, ACCOUNT_2],
},
async ({ controller }) => {
await controller.refresh();

expect(controller.state).toStrictEqual({
accounts: {
[CHECKSUM_ADDRESS_1]: { balance: '0x10', stakedBalance: '0x1' },
[CHECKSUM_ADDRESS_2]: { balance: '0x0' },
},
accountsByChainId: {
'0x1': {
[CHECKSUM_ADDRESS_1]: {
balance: '0x10',
stakedBalance: '0x1',
},
[CHECKSUM_ADDRESS_2]: {
balance: '0x0',
},
},
},
});
},
);
});

it('should not update staked balance when includeStakedAssets is disabled', async () => {
mockedQuery
.mockReturnValueOnce(Promise.resolve('0x13'))
.mockReturnValueOnce(Promise.resolve('0x14'));

await withController(
{
options: {
includeStakedAssets: false,
getStakedBalanceForChain: jest.fn().mockResolvedValue('0x1'),
},
isMultiAccountBalancesEnabled: false,
selectedAccount: ACCOUNT_1,
listAccounts: [ACCOUNT_1, ACCOUNT_2],
},
async ({ controller }) => {
await controller.refresh();

expect(controller.state).toStrictEqual({
accounts: {
[CHECKSUM_ADDRESS_1]: { balance: '0x13' },
[CHECKSUM_ADDRESS_2]: { balance: '0x0' },
},
accountsByChainId: {
'0x1': {
[CHECKSUM_ADDRESS_1]: {
balance: '0x13',
},
[CHECKSUM_ADDRESS_2]: {
balance: '0x0',
},
},
},
});
},
);
});

it('should update staked balance when includeStakedAssets and multi-account is enabled', async () => {
mockedQuery
.mockReturnValueOnce(Promise.resolve('0x11'))
.mockReturnValueOnce(Promise.resolve('0x12'));

await withController(
{
options: {
includeStakedAssets: true,
getStakedBalanceForChain: jest.fn().mockResolvedValue('0x1'),
},
isMultiAccountBalancesEnabled: true,
selectedAccount: ACCOUNT_1,
listAccounts: [ACCOUNT_1, ACCOUNT_2],
},
async ({ controller }) => {
await controller.refresh();

expect(controller.state).toStrictEqual({
accounts: {
[CHECKSUM_ADDRESS_1]: { balance: '0x11', stakedBalance: '0x1' },
[CHECKSUM_ADDRESS_2]: { balance: '0x12', stakedBalance: '0x1' },
},
accountsByChainId: {
'0x1': {
[CHECKSUM_ADDRESS_1]: {
balance: '0x11',
stakedBalance: '0x1',
},
[CHECKSUM_ADDRESS_2]: {
balance: '0x12',
stakedBalance: '0x1',
},
},
},
});
},
);
});
});

describe('with networkClientId', () => {
Expand Down Expand Up @@ -438,6 +555,185 @@ describe('AccountTrackerController', () => {
},
);
});

it('should update staked balance when includeStakedAssets is enabled', async () => {
const networkClientId = 'holesky';
mockedQuery
.mockReturnValueOnce(Promise.resolve('0x10'))
.mockReturnValueOnce(Promise.resolve('0x11'));

await withController(
{
options: {
includeStakedAssets: true,
getStakedBalanceForChain: jest.fn().mockResolvedValue('0x1'),
},
isMultiAccountBalancesEnabled: false,
selectedAccount: ACCOUNT_1,
listAccounts: [ACCOUNT_1, ACCOUNT_2],
networkClientById: {
[networkClientId]: buildCustomNetworkClientConfiguration({
chainId: '0x4268',
}),
},
},
async ({ controller }) => {
await controller.refresh();

expect(controller.state).toStrictEqual({
accounts: {
[CHECKSUM_ADDRESS_1]: { balance: '0x10', stakedBalance: '0x1' },
[CHECKSUM_ADDRESS_2]: { balance: '0x0' },
},
accountsByChainId: {
'0x1': {
[CHECKSUM_ADDRESS_1]: {
balance: '0x10',
stakedBalance: '0x1',
},
[CHECKSUM_ADDRESS_2]: {
balance: '0x0',
},
},
},
});
},
);
});

it('should not update staked balance when includeStakedAssets is disabled', async () => {
const networkClientId = 'holesky';
mockedQuery
.mockReturnValueOnce(Promise.resolve('0x13'))
.mockReturnValueOnce(Promise.resolve('0x14'));

await withController(
{
options: {
includeStakedAssets: false,
getStakedBalanceForChain: jest.fn().mockResolvedValue('0x1'),
},
isMultiAccountBalancesEnabled: false,
selectedAccount: ACCOUNT_1,
listAccounts: [ACCOUNT_1, ACCOUNT_2],
networkClientById: {
[networkClientId]: buildCustomNetworkClientConfiguration({
chainId: '0x4268',
}),
},
},
async ({ controller }) => {
await controller.refresh();

expect(controller.state).toStrictEqual({
accounts: {
[CHECKSUM_ADDRESS_1]: { balance: '0x13' },
[CHECKSUM_ADDRESS_2]: { balance: '0x0' },
},
accountsByChainId: {
'0x1': {
[CHECKSUM_ADDRESS_1]: {
balance: '0x13',
},
[CHECKSUM_ADDRESS_2]: {
balance: '0x0',
},
},
},
});
},
);
});

it('should update staked balance when includeStakedAssets and multi-account is enabled', async () => {
const networkClientId = 'holesky';
mockedQuery
.mockReturnValueOnce(Promise.resolve('0x11'))
.mockReturnValueOnce(Promise.resolve('0x12'));

await withController(
{
options: {
includeStakedAssets: true,
getStakedBalanceForChain: jest.fn().mockResolvedValue('0x1'),
},
isMultiAccountBalancesEnabled: true,
selectedAccount: ACCOUNT_1,
listAccounts: [ACCOUNT_1, ACCOUNT_2],
networkClientById: {
[networkClientId]: buildCustomNetworkClientConfiguration({
chainId: '0x4268',
}),
},
},
async ({ controller }) => {
await controller.refresh();

expect(controller.state).toStrictEqual({
accounts: {
[CHECKSUM_ADDRESS_1]: { balance: '0x11', stakedBalance: '0x1' },
[CHECKSUM_ADDRESS_2]: { balance: '0x12', stakedBalance: '0x1' },
},
accountsByChainId: {
'0x1': {
[CHECKSUM_ADDRESS_1]: {
balance: '0x11',
stakedBalance: '0x1',
},
[CHECKSUM_ADDRESS_2]: {
balance: '0x12',
stakedBalance: '0x1',
},
},
},
});
},
);
});

it('should not update staked balance when includeStakedAssets and multi-account is enabled if network unsupported', async () => {
const networkClientId = 'polygon';
mockedQuery
.mockReturnValueOnce(Promise.resolve('0x11'))
.mockReturnValueOnce(Promise.resolve('0x12'));

await withController(
{
options: {
includeStakedAssets: true,
getStakedBalanceForChain: jest.fn().mockResolvedValue(undefined),
},
isMultiAccountBalancesEnabled: true,
selectedAccount: ACCOUNT_1,
listAccounts: [ACCOUNT_1, ACCOUNT_2],
networkClientById: {
[networkClientId]: buildCustomNetworkClientConfiguration({
chainId: '0x89',
}),
},
},
async ({ controller }) => {
await controller.refresh();

expect(controller.state).toStrictEqual({
accounts: {
[CHECKSUM_ADDRESS_1]: { balance: '0x11' },
[CHECKSUM_ADDRESS_2]: { balance: '0x12' },
},
accountsByChainId: {
'0x1': {
[CHECKSUM_ADDRESS_1]: {
balance: '0x11',
},
[CHECKSUM_ADDRESS_2]: {
balance: '0x12',
},
},
},
});
},
);
});
});
});

Expand All @@ -462,6 +758,33 @@ describe('AccountTrackerController', () => {
},
);
});

it('should sync staked balance with addresses', async () => {
await withController(
{
options: {
includeStakedAssets: true,
getStakedBalanceForChain: jest.fn().mockResolvedValue('0x1'),
},
isMultiAccountBalancesEnabled: true,
selectedAccount: ACCOUNT_1,
listAccounts: [],
},
async ({ controller }) => {
mockedQuery
.mockReturnValueOnce(Promise.resolve('0x10'))
.mockReturnValueOnce(Promise.resolve('0x20'));
const result = await controller.syncBalanceWithAddresses([
ADDRESS_1,
ADDRESS_2,
]);
expect(result[ADDRESS_1].balance).toBe('0x10');
expect(result[ADDRESS_2].balance).toBe('0x20');
expect(result[ADDRESS_1].stakedBalance).toBe('0x1');
expect(result[ADDRESS_2].stakedBalance).toBe('0x1');
},
);
});
});

it('should call refresh every interval on legacy polling', async () => {
Expand Down Expand Up @@ -647,8 +970,8 @@ async function withController<ReturnValue>(

const controller = new AccountTrackerController({
messenger: accountTrackerMessenger,
...options,
getStakedBalanceForChain: jest.fn(),
...options,
});

return await testFunction({
Expand Down
15 changes: 9 additions & 6 deletions packages/assets-controllers/src/AccountTrackerController.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,10 @@ import { type Hex, assert } from '@metamask/utils';
import { Mutex } from 'async-mutex';
import { cloneDeep } from 'lodash';

import type { AssetsContractController } from './AssetsContractController';
import type {
AssetsContractController,
StakedBalance,
} from './AssetsContractController';

/**
* The name of the {@link AccountTrackerController}.
Expand Down Expand Up @@ -428,19 +431,19 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
async syncBalanceWithAddresses(
addresses: string[],
networkClientId?: NetworkClientId,
): Promise<Record<string, { balance: string }>> {
): Promise<
Record<string, { balance: string; stakedBalance?: StakedBalance }>
> {
const { ethQuery } = this.#getCorrectNetworkClient(networkClientId);

return await Promise.all(
addresses.map(
(
address,
): Promise<[string, string, string | undefined] | undefined> => {
(address): Promise<[string, string, StakedBalance] | undefined> => {
return safelyExecuteWithTimeout(async () => {
assert(ethQuery, 'Provider not set.');
const balance = await query(ethQuery, 'getBalance', [address]);

let stakedBalance: string | undefined;
let stakedBalance: StakedBalance;
if (this.#includeStakedAssets) {
stakedBalance = await this.#getStakedBalanceForChain(
address,
Expand Down
Loading

0 comments on commit cb9a55c

Please sign in to comment.