Skip to content

Commit

Permalink
refactor: switch alchemy client configs to use alchemy transport
Browse files Browse the repository at this point in the history
  • Loading branch information
moldy530 committed Sep 27, 2024
1 parent 3a78277 commit 9de5775
Show file tree
Hide file tree
Showing 37 changed files with 547 additions and 505 deletions.
7 changes: 2 additions & 5 deletions aa-sdk/core/src/client/smartAccountClient.test.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { custom, type Transaction } from "viem";
import { polygonMumbai } from "viem/chains";
import type { SpyInstance } from "vitest";
import type { MockInstance } from "vitest";
import * as receiptActions from "../actions/bundler/getUserOperationReceipt.js";
import type { UserOperationReceipt } from "../types.js";
import {
Expand Down Expand Up @@ -110,10 +110,7 @@ describe("SmartAccountClient Tests", async () => {
const thenExpectRetriesToBe = async (
expectedRetryMsDelays: number[],
expectedMockCalls: number,
getUserOperationReceiptMock: SpyInstance<
any,
Promise<UserOperationReceipt | null>
>
getUserOperationReceiptMock: MockInstance
) => {
expect(retryMsDelays).toEqual(expectedRetryMsDelays);
expect(getUserOperationReceiptMock).toHaveBeenCalledTimes(
Expand Down
8 changes: 8 additions & 0 deletions account-kit/core/src/actions/getAlchemyTransport.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
import type { AlchemyTransport } from "@account-kit/infra";
import type { AlchemyAccountsConfig } from "../types";

export function getAlchemyTransport(
config: AlchemyAccountsConfig
): AlchemyTransport {
return config.store.getState().transport;
}
45 changes: 18 additions & 27 deletions account-kit/core/src/actions/getSmartAccountClient.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import {
createAlchemySmartAccountClientFromExisting,
createAlchemySmartAccountClient,
type AlchemySmartAccountClient,
type AlchemySmartAccountClientConfig,
} from "@account-kit/infra";
Expand All @@ -19,7 +19,7 @@ import {
type MultiOwnerPluginActions,
type PluginManagerActions,
} from "@account-kit/smart-contracts";
import type { Address, Chain, Transport } from "viem";
import type { Address, Chain } from "viem";
import type {
AlchemyAccountsConfig,
Connection,
Expand All @@ -29,21 +29,16 @@ import type {
} from "../types";
import { createAccount } from "./createAccount.js";
import { getAccount, type GetAccountParams } from "./getAccount.js";
import { getBundlerClient } from "./getBundlerClient.js";
import { getAlchemyTransport } from "./getAlchemyTransport.js";
import { getConnection } from "./getConnection.js";
import { getSignerStatus } from "./getSignerStatus.js";

export type GetSmartAccountClientParams<
TTransport extends Transport = Transport,
TChain extends Chain | undefined = Chain | undefined,
TAccount extends SupportedAccountTypes = SupportedAccountTypes
> = Omit<
AlchemySmartAccountClientConfig<
TTransport,
TChain,
SupportedAccount<TAccount>
>,
"rpcUrl" | "chain" | "apiKey" | "jwt" | "account"
AlchemySmartAccountClientConfig<TChain, SupportedAccount<TAccount>>,
"transport" | "account" | "chain"
> &
GetAccountParams<TAccount>;

Expand All @@ -60,29 +55,22 @@ export type ClientActions<
: never;

export type GetSmartAccountClientResult<
TTransport extends Transport = Transport,
TChain extends Chain | undefined = Chain | undefined,
TAccount extends SupportedAccounts = SupportedAccounts
> = {
client?: AlchemySmartAccountClient<
TTransport,
TChain,
TAccount,
ClientActions<TAccount>
>;
client?: AlchemySmartAccountClient<TChain, TAccount, ClientActions<TAccount>>;
address?: Address;
isLoadingClient: boolean;
error?: Error;
};

export function getSmartAccountClient<
TTransport extends Transport = Transport,
TChain extends Chain | undefined = Chain | undefined,
TAccount extends SupportedAccountTypes = SupportedAccountTypes
>(
params: GetSmartAccountClientParams<TTransport, TChain, TAccount>,
params: GetSmartAccountClientParams<TChain, TAccount>,
config: AlchemyAccountsConfig
): GetSmartAccountClientResult<TTransport, TChain, SupportedAccount<TAccount>>;
): GetSmartAccountClientResult<TChain, SupportedAccount<TAccount>>;

/**
* Obtains a smart account client based on the provided parameters and configuration. Supports creating any of the SupportAccountTypes in Account Kit.
Expand Down Expand Up @@ -116,7 +104,7 @@ export function getSmartAccountClient(
config
);
const signerStatus = getSignerStatus(config);
const bundlerClient = getBundlerClient(config);
const transport = getAlchemyTransport(config);
const connection = getConnection(config);
const clientState =
config.store.getState().smartAccountClients[connection.chain.id]?.[type];
Expand Down Expand Up @@ -189,8 +177,9 @@ export function getSmartAccountClient(
switch (account.source) {
case "LightAccount":
return {
client: createAlchemySmartAccountClientFromExisting({
client: bundlerClient,
client: createAlchemySmartAccountClient({
transport,
chain: connection.chain,
account: account,
policyId: connection.policyId,
...clientParams,
Expand All @@ -200,8 +189,9 @@ export function getSmartAccountClient(
};
case "MultiOwnerLightAccount":
return {
client: createAlchemySmartAccountClientFromExisting({
client: bundlerClient,
client: createAlchemySmartAccountClient({
transport,
chain: connection.chain,
account: account,
policyId: connection.policyId,
...clientParams,
Expand All @@ -211,8 +201,9 @@ export function getSmartAccountClient(
};
case "MultiOwnerModularAccount":
return {
client: createAlchemySmartAccountClientFromExisting({
client: bundlerClient,
client: createAlchemySmartAccountClient({
transport,
chain: connection.chain,
account: account,
policyId: connection.policyId,
...clientParams,
Expand Down
6 changes: 4 additions & 2 deletions account-kit/core/src/actions/setChain.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { createAlchemyPublicRpcClient } from "@account-kit/infra";
import { alchemy, createAlchemyPublicRpcClient } from "@account-kit/infra";
import { switchChain } from "@wagmi/core";
import type { Chain } from "viem";
import { ChainNotFoundError } from "../errors.js";
Expand Down Expand Up @@ -27,12 +27,14 @@ export async function setChain(config: AlchemyAccountsConfig, chain: Chain) {
}

await switchChain(config._internal.wagmiConfig, { chainId: chain.id });
const transport = alchemy(connection);

config.store.setState(() => ({
chain,
transport,
bundlerClient: createAlchemyPublicRpcClient({
transport,
chain,
connectionConfig: connection,
}),
}));
}
2 changes: 2 additions & 0 deletions account-kit/core/src/actions/watchSmartAccountClient.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@ import { setChain } from "./setChain.js";
import { watchSmartAccountClient } from "./watchSmartAccountClient.js";

describe("watchSmartAccountClient", () => {
beforeEach(() => localStorage.clear());

it("should fire the on subscribe callback if signer status changes", () => {
const config = givenConfig();
const onChange = vi.fn();
Expand Down
23 changes: 10 additions & 13 deletions account-kit/core/src/actions/watchSmartAccountClient.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import type { Chain, Transport } from "viem";
import type { Chain } from "viem";
import { ClientOnlyPropertyError } from "../errors.js";
import type {
AlchemyAccountsConfig,
Expand Down Expand Up @@ -27,28 +27,23 @@ import {
* @template TTransport extends Transport = Transport
* @template TChain extends Chain | undefined = Chain | undefined
*
* @param {GetSmartAccountClientParams<TTransport, TChain, TAccount>} params the parameters needed to get the smart account client
* @param {GetSmartAccountClientParams<TChain, TAccount>} params the parameters needed to get the smart account client
* @param {AlchemyAccountsConfig} config the configuration containing the client store and other settings
* @returns {(onChange: (client: GetSmartAccountClientResult<TTransport, TChain, SupportedAccount<TAccount>>) => void) => (() => void)} a function that accepts a callback to be called when the client changes and returns a function to unsubscribe from the store
* @returns {(onChange: (client: GetSmartAccountClientResult<TChain, SupportedAccount<TAccount>>) => void) => (() => void)} a function that accepts a callback to be called when the client changes and returns a function to unsubscribe from the store
*/ export function watchSmartAccountClient<
TAccount extends SupportedAccountTypes,
TTransport extends Transport = Transport,
TChain extends Chain | undefined = Chain | undefined
>(
params: GetSmartAccountClientParams<TTransport, TChain, TAccount>,
params: GetSmartAccountClientParams<TChain, TAccount>,
config: AlchemyAccountsConfig
) {
return (
onChange: (
client: GetSmartAccountClientResult<
TTransport,
TChain,
SupportedAccount<TAccount>
>
client: GetSmartAccountClientResult<TChain, SupportedAccount<TAccount>>
) => void
) => {
const accounts = config.store.getState().accounts;
if (!accounts) {
const accounts_ = config.store.getState().accounts;
if (!accounts_) {
throw new ClientOnlyPropertyError("account");
}

Expand All @@ -58,7 +53,9 @@ import {
account: accounts![chain.id][params.type],
bundlerClient,
}),
() => onChange(getSmartAccountClient(params, config)),
() => {
onChange(getSmartAccountClient(params, config));
},
{
equalityFn(a, b) {
return (
Expand Down
1 change: 0 additions & 1 deletion account-kit/core/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,6 @@ export {
export type * from "@account-kit/infra";
export {
createAlchemySmartAccountClient,
createAlchemySmartAccountClientFromExisting,
type AlchemySmartAccountClient,
type AlchemySmartAccountClientConfig,
} from "@account-kit/infra";
80 changes: 80 additions & 0 deletions account-kit/core/src/store/store.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
import { arbitrumSepolia, sepolia } from "@account-kit/infra";
import { setChain } from "../actions/setChain.js";
import { createConfig } from "../createConfig.js";
import { createDefaultAccountState } from "./store.js";
import { DEFAULT_STORAGE_KEY } from "./types.js";

describe("createConfig tests", () => {
beforeEach(() => localStorage.clear());

it("should setup the config with the correct transport", () => {
const config = givenConfig();

expect({ ...config.store.getState().transport }).toMatchInlineSnapshot(`
{
"config": {
"policyId": undefined,
"rpcUrl": "/api/sepolia",
},
"updateHeaders": [Function],
}
`);
});

it("should omit the transport when persisting to storage", () => {
givenConfig();

expect(getStorageItem("transport")).toMatchInlineSnapshot(`undefined`);
});

it("should rehydrate the transport based on the bundler client", async () => {
const config = givenConfig();

// update the chain so we can make sure the store is updated
expect(getStorageItem("bundlerClient").connection.chain.id).toBe(
sepolia.id
);
await setChain(config, arbitrumSepolia);
expect(getStorageItem("bundlerClient").connection.chain.id).toBe(
arbitrumSepolia.id
);

// create a config that is the result of a rehydration
const hydratedConfig = givenConfig();
await hydratedConfig.store.persist.rehydrate();
expect(hydratedConfig.store.getState().transport.config)
.toMatchInlineSnapshot(`
{
"rpcUrl": "/api/arbitrumSepolia",
}
`);

expect(hydratedConfig.store.getState().bundlerClient.chain.id).toBe(
config.store.getState().bundlerClient.chain.id
);
});

const getStorageItem = (name: string) => {
return JSON.parse(localStorage.getItem(DEFAULT_STORAGE_KEY) ?? "{}")[
"state"
][name];
};

const givenConfig = () => {
const config = createConfig({
chain: sepolia,
connections: [
{ chain: sepolia, rpcUrl: "/api/sepolia" },
{ chain: arbitrumSepolia, rpcUrl: "/api/arbitrumSepolia" },
],
signerConnection: { rpcUrl: "/api/signer" },
storage: () => localStorage,
});

config.store.setState({
accounts: createDefaultAccountState([sepolia, arbitrumSepolia]),
});

return config;
};
});
Loading

0 comments on commit 9de5775

Please sign in to comment.