Skip to content

Commit 2f4a32f

Browse files
AnsonhkgAnsonhkg
committed
refactor(artillery): use master account instead of generated account
- Introduced AccountManager for handling master account operations, including fetching account details and managing payment balances. - Added NetworkManager to streamline network module initialization and public client creation. - Updated init script to utilize new managers for improved clarity and maintainability. - Refactored pkpSign processor to support master account authentication and context creation. - Adjusted package.json to simplify artillery initialization command.
1 parent f0af2cd commit 2f4a32f

File tree

8 files changed

+249
-180
lines changed

8 files changed

+249
-180
lines changed

‎e2e/artillery/src/AccountManager.ts‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
import * as NetworkManager from "../../../e2e/src/helper/NetworkManager";
2+
import { privateKeyToAccount } from "viem/accounts";
3+
import { createPublicClient, formatEther, http } from "viem";
4+
import { createLitClient } from "@lit-protocol/lit-client";
5+
import * as StateManager from "./StateManager";
6+
import { printAligned } from "../../../e2e/src/helper/utils";
7+
8+
export const getMasterAccount = async () => {
9+
const privateKey = process.env['LIVE_MASTER_ACCOUNT'];
10+
11+
if (!privateKey) {
12+
throw new Error('❌ LIVE_MASTER_ACCOUNT environment variable is required');
13+
}
14+
15+
return privateKeyToAccount(privateKey as `0x${string}`);
16+
}
17+
18+
export const getAccountDetails = async ({
19+
account,
20+
publicClient,
21+
litClient,
22+
accountLabel = "Account",
23+
}: {
24+
account: any;
25+
publicClient: any;
26+
litClient: Awaited<ReturnType<typeof createLitClient>>;
27+
accountLabel?: string;
28+
}) => {
29+
console.log(`\n========== ${accountLabel} Details ==========`);
30+
31+
// Get all the data first
32+
const ethBalance = formatEther(await publicClient.getBalance({
33+
address: account.address,
34+
}));
35+
36+
const paymentManager = await litClient.getPaymentManager({
37+
account: account,
38+
});
39+
40+
const paymentBalance = await paymentManager.getBalance({
41+
userAddress: account.address
42+
});
43+
44+
// Determine status
45+
let statusLabel = "";
46+
let statusValue = "";
47+
48+
if (Number(paymentBalance.availableBalance) < 0) {
49+
statusLabel = "🚨 Status:";
50+
statusValue = `Negative balance (debt): ${paymentBalance.availableBalance}`;
51+
}
52+
53+
// Print all information with consistent alignment
54+
printAligned([
55+
{ label: "🔑 Address:", value: account.address },
56+
{ label: "💰 ETH Balance:", value: `${ethBalance} ETH` },
57+
{ label: "💳 Ledger Total Balance:", value: paymentBalance.totalBalance },
58+
{ label: "💳 Ledger Available Balance:", value: paymentBalance.availableBalance },
59+
{ label: statusLabel, value: statusValue },
60+
]);
61+
62+
return {
63+
ethBalance,
64+
ledgerBalance: paymentBalance.availableBalance,
65+
paymentManager,
66+
}
67+
}
68+

‎e2e/artillery/src/StateManager.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,10 @@ const FILE_NAME = 'artillery-state.json';
55

66
// State Object
77
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,
8+
masterAccount: {
9+
// privateKey: undefined as string | `0x${string}` | undefined,
10+
// address: undefined as string | `0x${string}` | undefined,
11+
authData: undefined as any | undefined, // Changed from string to any since authData is an object
1212
pkp: undefined as any | undefined,
1313
},
1414
};

‎e2e/artillery/src/init.ts‎

Lines changed: 77 additions & 145 deletions
Original file line numberDiff line numberDiff line change
@@ -1,121 +1,65 @@
1-
import { generatePrivateKey, privateKeyToAccount } from "viem/accounts";
2-
import { nagaDev } from "@lit-protocol/networks";
3-
import { createPublicClient, formatEther, http, parseEther } from "viem";
1+
import '../../src/helper/supressLogs';
42
import { createAuthManager, storagePlugins, ViemAccountAuthenticator } from "@lit-protocol/auth";
5-
import { fundAccount } from "../../../e2e/src/helper/fundAccount";
63
import * as StateManager from "./StateManager";
74
import { createLitClient } from "@lit-protocol/lit-client";
85
import { getOrCreatePkp } from "../../../e2e/src/helper/pkp-utils";
6+
import * as NetworkManager from "../../../e2e/src/helper/NetworkManager";
7+
import * as AccountManager from '../src/AccountManager';
98

109
const _network = process.env['NETWORK'];
1110

1211
// CONFIGURATIONS
1312
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-
}
13+
const LEDGER_MINIMUM_BALANCE = 10000;
2914

3015
(async () => {
3116

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-
4317
// -- Start
4418
console.log("\x1b[90m✅ Initialising Artillery...\x1b[0m");
4519

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-
});
20+
// 1. Setup network and chain client
21+
const networkModule = await NetworkManager.getLitNetworkModule();
22+
const publicClient = await NetworkManager.getViemPublicClient({ networkModule });
23+
const litClient = await createLitClient({ network: networkModule });
8424

85-
console.log("\x1b[90m✅ Creating Lit Client...\x1b[0m");
86-
const litClient = await createLitClient({ network: _networkModule });
25+
// 2. Setup the master account
26+
const masterAccount = await AccountManager.getMasterAccount();
8727

88-
console.log("\x1b[90m✅ Getting Live Master Account Payment Manager...\x1b[0m");
89-
const masterPaymentManager = await litClient.getPaymentManager({
90-
account: liveMasterAccount,
28+
const masterAccountDetails = await AccountManager.getAccountDetails({
29+
accountLabel: "Master Account",
30+
account: masterAccount,
31+
publicClient,
32+
litClient,
9133
});
9234

93-
// Deposit
94-
const masterPaymentBalance = await masterPaymentManager.getBalance({ userAddress: liveMasterAccount.address })
95-
console.log('✅ Live Master Account Payment Balance:', masterPaymentBalance);
35+
if (Number(masterAccountDetails.ethBalance) < REJECT_BALANCE_THRESHOLD) {
36+
throw new Error(`🚨 Live Master Account Balance is less than REJECT_BALANCE_THRESHOLD: ${REJECT_BALANCE_THRESHOLD} ETH`);
37+
}
9638

97-
if (LEDGER_MINIMUM_BALANCE > Number(masterPaymentBalance.availableBalance)) {
39+
if (LEDGER_MINIMUM_BALANCE > Number(masterAccountDetails.ledgerBalance)) {
9840

9941
// 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);
42+
const difference = LEDGER_MINIMUM_BALANCE - Number(masterAccountDetails.ledgerBalance);
43+
44+
console.log(`🚨 Live Master Account Ledger Balance is less than LEDGER_MINIMUM_BALANCE: ${LEDGER_MINIMUM_BALANCE} ETH. Attempting to top up the difference of ${difference} ETH to the master account.`);
10245

10346
// deposit the difference
10447
console.log("\x1b[90m✅ Depositing the difference to Live Master Account Payment Manager...\x1b[0m");
105-
await masterPaymentManager.deposit({ amountInEth: difference.toString() });
48+
await masterAccountDetails.paymentManager.deposit({ amountInEth: difference.toString() });
10649

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');
50+
// print the new balance
51+
const newBalance = await masterAccountDetails.paymentManager.getBalance({ userAddress: masterAccount.address });
52+
console.log('✅ New Live Master Account Payment Balance:', newBalance.availableBalance);
11253
}
11354

114-
/**
115-
* ====================================
116-
* Initialise the AuthManager
117-
* ====================================
118-
*/
55+
// 3. Authenticate the master account and store the auth data
56+
const masterAccountAuthData = await StateManager.getOrUpdate(
57+
'masterAccount.authData',
58+
await ViemAccountAuthenticator.authenticate(masterAccount),
59+
);
60+
console.log('✅ Master Account Auth Data:', masterAccountAuthData);
61+
62+
// 4. initialise the auth manager
11963
const authManager = createAuthManager({
12064
storage: storagePlugins.localStorageNode({
12165
appName: 'artillery-testing-app',
@@ -124,76 +68,64 @@ if (!config) {
12468
}),
12569
});
12670

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-
},
71+
// 5. get or mint a PKP for the master account
72+
const masterAccountPkp = await StateManager.getOrUpdate(
73+
'masterAccount.pkp',
74+
await getOrCreatePkp(
75+
litClient,
76+
masterAccountAuthData,
77+
masterAccount,
78+
'./artillery-pkp-tokens',
79+
`${_network}-artillery`
80+
)
81+
);
82+
83+
console.log('✅ Master Account PKP:', masterAccountPkp);
84+
85+
// create pkp auth context
86+
const masterAccountPkpAuthContext = await authManager.createPkpAuthContext({
87+
authData: masterAccountAuthData,
88+
pkpPublicKey: masterAccountPkp.publicKey,
13789
authConfig: {
138-
statement: 'I authorize the Lit Protocol to execute this Lit Action.',
139-
domain: 'example.com',
14090
resources: [
141-
['lit-action-execution', '*'],
14291
['pkp-signing', '*'],
92+
['lit-action-execution', '*'],
14393
['access-control-condition-decryption', '*'],
14494
],
145-
capabilityAuthSigs: [],
14695
expiration: new Date(Date.now() + 1000 * 60 * 15).toISOString(),
14796
},
14897
litClient: litClient,
14998
});
15099

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-
}
100+
console.log('✅ Master Account PKP Auth Context:', masterAccountPkpAuthContext);
101+
102+
// 6. create the auth context (this should be generated each time)
103+
// const masterAccountAuthContext = await authManager.createEoaAuthContext({
104+
// config: {
105+
// account: masterAccount,
106+
// },
107+
// authConfig: {
108+
// statement: 'I authorize the Lit Protocol to execute this Lit Action.',
109+
// domain: 'example.com',
110+
// resources: [
111+
// ['lit-action-execution', '*'],
112+
// ['pkp-signing', '*'],
113+
// ['access-control-condition-decryption', '*'],
114+
// ],
115+
// capabilityAuthSigs: [],
116+
// expiration: new Date(Date.now() + 1000 * 60 * 15).toISOString(),
117+
// },
118+
// litClient: litClient,
119+
// });
188120

189-
// run pkpSign test
190-
// console.log("\x1b[90m✅ Running PKP Sign Test...\x1b[0m");
121+
// (uncomment to test) run the pkpSign endpoint
191122
// const res = await litClient.chain.ethereum.pkpSign({
192-
// authContext: aliceEoaAuthContext,
193-
// pubKey: aliceViemAccountPkp.publicKey,
123+
// authContext: masterAccountAuthContext,
124+
// pubKey: masterAccountPkp.publicKey,
194125
// toSign: 'Hello, world!',
195126
// });
196127

197128
// console.log('✅ PKP Sign Test Result:', res);
129+
198130
process.exit();
199131
})();

0 commit comments

Comments
 (0)