Skip to content

Commit 68ad89d

Browse files
AnsonhkgAnsonhkg
committed
feat(artillery): add PKP sign test and state management for E2E testing
- Introduced new artillery configurations for PKP sign testing. - Implemented state management in StateManager to handle account data. - Created init script to initialise necessary resources and manage account funding. - Added pkpSign processor to execute signing operations during tests. - Updated package.json to include new artillery commands.
1 parent 4970ac4 commit 68ad89d

File tree

5 files changed

+495
-6
lines changed

5 files changed

+495
-6
lines changed

‎e2e/artillery/configs/pkp-sign.yml‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
config:
2+
target: "dummy"
3+
phases:
4+
- duration: 30
5+
arrivalRate: 1
6+
maxVusers: 3
7+
name: "PKP Sign Test"
8+
processor: "../src/processors/pkpSign.ts"
9+
environments:
10+
production:
11+
target: "dummy"
12+
13+
scenarios:
14+
- name: "PKP Sign Load Test"
15+
weight: 100
16+
flow:
17+
- function: "runPkpSignTest"
18+
- think: 2

‎e2e/artillery/src/StateManager.ts‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
import fs from 'fs/promises';
2+
3+
// CONFIGURATIONS
4+
const FILE_NAME = 'artillery-state.json';
5+
6+
// State Object
7+
const StateObject = {
8+
aliceViemEoaAccount: {
9+
privateKey: undefined as string | `0x${string}` | undefined,
10+
address: undefined as string | `0x${string}` | undefined,
11+
authData: undefined as string | undefined,
12+
pkp: undefined as any | undefined,
13+
},
14+
};
15+
16+
// STATE - derived from StateObject to ensure type consistency
17+
type State = typeof StateObject;
18+
19+
// read the file if it exists, if not throw an error
20+
export const readFile = async (): Promise<State> => {
21+
async function _readFile() {
22+
const file = await fs.readFile(FILE_NAME, 'utf8');
23+
const content = JSON.parse(file) as State;
24+
25+
// If content is empty object, write base state
26+
if (Object.keys(content).length === 0) {
27+
await fs.writeFile(FILE_NAME, JSON.stringify(StateObject, null, 2));
28+
return StateObject;
29+
}
30+
31+
return content;
32+
}
33+
34+
try {
35+
return await _readFile();
36+
} catch (error) {
37+
console.log("🚨 Failed to read file, creating new file...");
38+
await createFile();
39+
return await _readFile();
40+
}
41+
};
42+
43+
// create the file if it doesn't exist
44+
export const createFile = async () => {
45+
await fs.writeFile(FILE_NAME, JSON.stringify(StateObject, null, 2));
46+
};
47+
48+
// Type-safe field paths - dynamically derived from State type
49+
type StatePaths = {
50+
[K in keyof State]: K extends string
51+
? State[K] extends object
52+
? {
53+
[P in keyof State[K]]: P extends string ? `${K}.${P}` : never
54+
}[keyof State[K]] | K // Include both nested paths AND top-level key
55+
: K
56+
: never
57+
}[keyof State];
58+
59+
// Helper type to get nested property type
60+
type GetNestedType<T, P extends string> =
61+
P extends `${infer K}.${infer Rest}`
62+
? K extends keyof T
63+
? Rest extends keyof T[K]
64+
? T[K][Rest]
65+
: never
66+
: never
67+
: P extends keyof T
68+
? T[P]
69+
: never;
70+
71+
// Map paths to their corresponding types
72+
type StatePathValue<T extends StatePaths> = GetNestedType<State, T>;
73+
74+
/**
75+
* Updates a specific field in the state with type safety
76+
* @param path - The dot-notation path to the field to update
77+
* @param value - The value to set, must match the field's type
78+
*/
79+
export const updateField = async <T extends StatePaths>(
80+
path: T,
81+
value: StatePathValue<T>
82+
): Promise<void> => {
83+
const state = await readFile();
84+
85+
// Split the path and navigate to the nested property
86+
const pathParts = path.split('.') as [keyof State, string];
87+
const [rootKey, nestedKey] = pathParts;
88+
89+
if (rootKey in state && typeof state[rootKey] === 'object' && state[rootKey] !== null) {
90+
(state[rootKey] as any)[nestedKey] = value;
91+
await fs.writeFile(FILE_NAME, JSON.stringify(state, null, 2));
92+
} else {
93+
throw new Error(`Invalid path: ${path}`);
94+
}
95+
};
96+
97+
/**
98+
* Gets a field value, or updates it if it doesn't exist (is undefined/null/empty string)
99+
* @param path - The dot-notation path to the field to get/update (or top-level key)
100+
* @param defaultValue - The value to set if the current value is undefined/null/empty
101+
* @returns The existing value or the newly set default value
102+
*/
103+
export const getOrUpdate = async <T extends StatePaths>(
104+
path: T,
105+
defaultValue: NonNullable<StatePathValue<T>>
106+
): Promise<NonNullable<StatePathValue<T>>> => {
107+
const state = await readFile();
108+
109+
// Check if it's a top-level property or nested property
110+
if (!path.includes('.')) {
111+
// Top-level property
112+
const currentValue = (state as any)[path];
113+
114+
// If value exists and is not null/undefined/empty string, return it
115+
if (currentValue != null && currentValue !== '') {
116+
return currentValue as NonNullable<StatePathValue<T>>;
117+
}
118+
119+
// Otherwise, update with default value and return it
120+
(state as any)[path] = defaultValue;
121+
await fs.writeFile(FILE_NAME, JSON.stringify(state, null, 2));
122+
return defaultValue;
123+
} else {
124+
// Nested property
125+
const pathParts = path.split('.') as [keyof State, string];
126+
const [rootKey, nestedKey] = pathParts;
127+
128+
if (rootKey in state && typeof state[rootKey] === 'object' && state[rootKey] !== null) {
129+
const currentValue = (state[rootKey] as any)[nestedKey];
130+
131+
// If value exists and is not null/undefined/empty string, return it
132+
if (currentValue != null && currentValue !== '') {
133+
return currentValue as NonNullable<StatePathValue<T>>;
134+
}
135+
136+
// Otherwise, update with default value and return it
137+
(state[rootKey] as any)[nestedKey] = defaultValue;
138+
await fs.writeFile(FILE_NAME, JSON.stringify(state, null, 2));
139+
return defaultValue;
140+
} else {
141+
throw new Error(`Invalid path: ${path}`);
142+
}
143+
}
144+
};

‎e2e/artillery/src/init.ts‎

Lines changed: 199 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,199 @@
1+
import { generatePrivateKey, privateKeyToAccount } from "viem/accounts";
2+
import { nagaDev } from "@lit-protocol/networks";
3+
import { createPublicClient, formatEther, http, parseEther } from "viem";
4+
import { createAuthManager, storagePlugins, ViemAccountAuthenticator } from "@lit-protocol/auth";
5+
import { fundAccount } from "../../../e2e/src/helper/fundAccount";
6+
import * as StateManager from "./StateManager";
7+
import { createLitClient } from "@lit-protocol/lit-client";
8+
import { getOrCreatePkp } from "../../../e2e/src/helper/pkp-utils";
9+
10+
const _network = process.env['NETWORK'];
11+
12+
// CONFIGURATIONS
13+
const REJECT_BALANCE_THRESHOLD = 0;
14+
const FUNDING_IF_LESS_THAN = 0.1;
15+
const FUNDING_AMOUNT = 0.1;
16+
const LEDGER_MINIMUM_BALANCE = 500;
17+
18+
const NETWORK_CONFIG = {
19+
'naga-dev': { importName: 'nagaDev', type: 'live' },
20+
'naga-test': { importName: 'nagaTest', type: 'live' },
21+
'naga-local': { importName: 'nagaLocal', type: 'local' },
22+
'naga-staging': { importName: 'nagaStaging', type: 'live' },
23+
} as const;
24+
25+
const config = NETWORK_CONFIG[_network as keyof typeof NETWORK_CONFIG];
26+
if (!config) {
27+
throw new Error(`❌ Invalid network: ${_network}`);
28+
}
29+
30+
(async () => {
31+
32+
// -- Setup
33+
const networksModule = await import('@lit-protocol/networks');
34+
const _networkModule = networksModule[config.importName];
35+
36+
const viemChainConfig = _networkModule.getChainConfig();
37+
38+
const publicClient = createPublicClient({
39+
chain: viemChainConfig,
40+
transport: http(),
41+
});
42+
43+
// -- Start
44+
console.log("\x1b[90m✅ Initialising Artillery...\x1b[0m");
45+
46+
// 1. Setup the master account (only for Live network for load testing))
47+
const liveMasterAccount = privateKeyToAccount(process.env['LIVE_MASTER_ACCOUNT'] as `0x${string}`);
48+
console.log("🔑 Live Master Account:", liveMasterAccount.address);
49+
50+
const balance = formatEther(await publicClient.getBalance({
51+
address: liveMasterAccount.address,
52+
}));
53+
54+
if (Number(balance) < REJECT_BALANCE_THRESHOLD) {
55+
throw new Error(`🚨 Live Master Account Balance is less than REJECT_BALANCE_THRESHOLD: ${REJECT_BALANCE_THRESHOLD} ETH`);
56+
}
57+
58+
console.log("💰 Live Master Account Balance:", balance, "ETH");
59+
60+
// 2. Create an PKP Auth EOA account
61+
console.log("\x1b[90m✅ Creating PKP EOA Account...\x1b[0m");
62+
const aliceEoaPrivateKey = await StateManager.getOrUpdate(
63+
'aliceViemEoaAccount.privateKey',
64+
generatePrivateKey()
65+
) as `0x${string}`;
66+
67+
console.log("🔑 Alice Viem EOA Account Private Key:", aliceEoaPrivateKey);
68+
69+
const aliceViemEoaAccount = privateKeyToAccount(aliceEoaPrivateKey);
70+
const aliceViemEoaAccountAuthData = await StateManager.getOrUpdate(
71+
'aliceViemEoaAccount.authData',
72+
JSON.stringify(await ViemAccountAuthenticator.authenticate(aliceViemEoaAccount))
73+
);
74+
75+
console.log("🔑 Alice Viem EOA Account:", aliceViemEoaAccount.address);
76+
console.log("🔑 Alice Viem EOA Account Auth Data:", aliceViemEoaAccountAuthData);
77+
78+
// 3. Fund the PKP Auth EOA account
79+
console.log("\x1b[90m✅ Funding PKP Auth EOA Account...\x1b[0m");
80+
await fundAccount(aliceViemEoaAccount, liveMasterAccount, _networkModule, {
81+
ifLessThan: FUNDING_IF_LESS_THAN.toString(),
82+
thenFundWith: FUNDING_AMOUNT.toString(),
83+
});
84+
85+
console.log("\x1b[90m✅ Creating Lit Client...\x1b[0m");
86+
const litClient = await createLitClient({ network: _networkModule });
87+
88+
console.log("\x1b[90m✅ Getting Live Master Account Payment Manager...\x1b[0m");
89+
const masterPaymentManager = await litClient.getPaymentManager({
90+
account: liveMasterAccount,
91+
});
92+
93+
// Deposit
94+
const masterPaymentBalance = await masterPaymentManager.getBalance({ userAddress: liveMasterAccount.address })
95+
console.log('✅ Live Master Account Payment Balance:', masterPaymentBalance);
96+
97+
if (LEDGER_MINIMUM_BALANCE > Number(masterPaymentBalance.availableBalance)) {
98+
99+
// find the difference between the minimum balance and the current balance
100+
const difference = LEDGER_MINIMUM_BALANCE - Number(masterPaymentBalance.availableBalance);
101+
console.log('💰 Difference:', difference);
102+
103+
// deposit the difference
104+
console.log("\x1b[90m✅ Depositing the difference to Live Master Account Payment Manager...\x1b[0m");
105+
await masterPaymentManager.deposit({ amountInEth: difference.toString() });
106+
107+
// get the new balance
108+
const newBalance = await masterPaymentManager.getBalance({ userAddress: liveMasterAccount.address })
109+
console.log('✅ New Live Master Account Payment Balance:', newBalance);
110+
} else {
111+
console.log('🔥 Live Master Account Payment Balance is greater than the minimum balance');
112+
}
113+
114+
/**
115+
* ====================================
116+
* Initialise the AuthManager
117+
* ====================================
118+
*/
119+
const authManager = createAuthManager({
120+
storage: storagePlugins.localStorageNode({
121+
appName: 'artillery-testing-app',
122+
networkName: `${_network}-artillery`,
123+
storagePath: './lit-auth-artillery',
124+
}),
125+
});
126+
127+
/**
128+
* ====================================
129+
* Create the auth context (recreate each time since it contains functions)
130+
* ====================================
131+
*/
132+
console.log("\x1b[90m✅ Creating Alice EOA Auth Context...\x1b[0m");
133+
const aliceEoaAuthContext = await authManager.createEoaAuthContext({
134+
config: {
135+
account: aliceViemEoaAccount,
136+
},
137+
authConfig: {
138+
statement: 'I authorize the Lit Protocol to execute this Lit Action.',
139+
domain: 'example.com',
140+
resources: [
141+
['lit-action-execution', '*'],
142+
['pkp-signing', '*'],
143+
['access-control-condition-decryption', '*'],
144+
],
145+
capabilityAuthSigs: [],
146+
expiration: new Date(Date.now() + 1000 * 60 * 15).toISOString(),
147+
},
148+
litClient: litClient,
149+
});
150+
151+
console.log('✅ Alice EOA Auth Context created');
152+
153+
// Get or create a PKP for alice
154+
console.log("\x1b[90m✅ Getting or creating PKP for Alice...\x1b[0m");
155+
const aliceViemAccountPkp = await StateManager.getOrUpdate(
156+
'aliceViemEoaAccount.pkp',
157+
await getOrCreatePkp(
158+
litClient,
159+
aliceViemEoaAccountAuthData,
160+
aliceViemEoaAccount,
161+
'./pkp-tokens',
162+
`${_network}-artillery`
163+
)
164+
);
165+
166+
console.log('✅ Alice Viem Account PKP:', aliceViemAccountPkp);
167+
168+
// Deposit for the aliceViemEoaAccount
169+
const aliceEoaViemAccountPaymentBalance = await masterPaymentManager.getBalance({ userAddress: aliceViemEoaAccount.address })
170+
console.log('✅ Alice EOA Viem Account Payment Balance:', aliceEoaViemAccountPaymentBalance);
171+
172+
if (LEDGER_MINIMUM_BALANCE > Number(aliceEoaViemAccountPaymentBalance.availableBalance)) {
173+
174+
// find the difference between the minimum balance and the current balance
175+
const difference = LEDGER_MINIMUM_BALANCE - Number(aliceEoaViemAccountPaymentBalance.availableBalance);
176+
console.log('💰 Difference:', difference);
177+
178+
// deposit the difference
179+
console.log("\x1b[90m✅ Depositing the difference to Alice EOA Viem Account Payment Manager...\x1b[0m");
180+
await masterPaymentManager.depositForUser({ userAddress: aliceViemEoaAccount.address, amountInEth: difference.toString() });
181+
182+
// get the new balance
183+
const newBalance = await masterPaymentManager.getBalance({ userAddress: aliceViemEoaAccount.address })
184+
console.log('✅ New Alice EOA Viem Account Payment Balance:', newBalance);
185+
} else {
186+
console.log('🔥 Alice EOA Viem Account Payment Balance is greater than the minimum balance');
187+
}
188+
189+
// run pkpSign test
190+
// console.log("\x1b[90m✅ Running PKP Sign Test...\x1b[0m");
191+
// const res = await litClient.chain.ethereum.pkpSign({
192+
// authContext: aliceEoaAuthContext,
193+
// pubKey: aliceViemAccountPkp.publicKey,
194+
// toSign: 'Hello, world!',
195+
// });
196+
197+
// console.log('✅ PKP Sign Test Result:', res);
198+
process.exit();
199+
})();

0 commit comments

Comments
 (0)