Skip to content

Commit 23bec70

Browse files
AnsonhkgAnsonhkg
committed
feat(e2e): add wrapped-keys test
1 parent 9f49188 commit 23bec70

File tree

5 files changed

+602
-8
lines changed

5 files changed

+602
-8
lines changed

‎packages/e2e/src/e2e.spec.ts‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ import {
2020
} from './helper/tests';
2121
import { init } from './init';
2222
import { AuthContext } from './types';
23-
import { createPregenDelegationServerReuseTest } from './tests/signSessionKey/pregen-delegation';
23+
import { createPregenDelegationServerReuseTest } from './test-helpers/signSessionKey/pregen-delegation';
2424

2525
const RPC_OVERRIDE = process.env['LIT_YELLOWSTONE_PRIVATE_RPC_URL'];
2626
if (RPC_OVERRIDE) {
@@ -276,7 +276,6 @@ describe('all', () => {
276276
authData: ctx.aliceViemAccountAuthData,
277277
pkpPublicKey: ctx.aliceViemAccountPkp.pubkey,
278278
clientLitClient: ctx.litClient,
279-
fallbackLitClient: ctx.litClient,
280279
resolvedNetwork: ctx.resolvedNetwork,
281280
})());
282281
});

‎packages/e2e/src/helper/fundAccount.ts‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -85,12 +85,19 @@ export const fundAccount = async (
8585
ifLessThan?: string;
8686
thenFundWith?: string;
8787
}
88-
) => {
88+
): Promise<
89+
| {
90+
txHash: `0x${string}`;
91+
rpcUrl: string;
92+
}
93+
| undefined
94+
> => {
8995
const defaultRpcUrl = networkModule.getChainConfig().rpcUrls.default.http[0];
9096
const isLocalNetwork = defaultRpcUrl.includes('127.0.0.1');
9197
const customRpcUrl = isLocalNetwork
9298
? process.env['LOCAL_RPC_URL']
9399
: process.env['LIT_YELLOWSTONE_PRIVATE_RPC_URL'];
100+
const rpcUrl = customRpcUrl || defaultRpcUrl;
94101

95102
if (customRpcUrl) {
96103
console.log(`🔧 Using custom E2E RPC URL: ***${customRpcUrl.slice(-6)}`);
@@ -103,7 +110,7 @@ export const fundAccount = async (
103110
// check account balance
104111
const publicClient = createPublicClient({
105112
chain: networkModule.getChainConfig(),
106-
transport: http(customRpcUrl || defaultRpcUrl),
113+
transport: http(rpcUrl),
107114
});
108115

109116
const balance = await publicClient.getBalance({
@@ -116,7 +123,7 @@ export const fundAccount = async (
116123

117124
const walletClient = createWalletClient({
118125
account: sponsorAccount,
119-
transport: http(customRpcUrl || defaultRpcUrl),
126+
transport: http(rpcUrl),
120127
});
121128

122129
// Get the next managed nonce for this sponsor account
@@ -130,14 +137,17 @@ export const fundAccount = async (
130137
account: sponsorAccount, // Add account for retry logic
131138
};
132139

133-
await sendTransactionWithRetry(
140+
const txHash = (await sendTransactionWithRetry(
134141
walletClient,
135142
transactionRequest,
136143
publicClient
137-
);
144+
)) as `0x${string}`;
138145

139146
console.log('✅ Topped up account with', options?.thenFundWith, 'ETH');
147+
return { txHash, rpcUrl };
140148
} else {
141149
console.log('✅ Account has enough balance');
142150
}
151+
152+
return undefined;
143153
};

‎packages/e2e/src/test-helpers/executeJs/wrappedKeys.ts‎

Lines changed: 205 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,205 @@
1+
import { randomBytes, createHash } from 'crypto';
2+
import { Wallet } from 'ethers';
3+
4+
import {
5+
api as wrappedKeysApi,
6+
constants as wrappedKeysConstants,
7+
} from '@lit-protocol/wrapped-keys';
8+
import type {
9+
LitClient,
10+
ExportPrivateKeyResult,
11+
GeneratePrivateKeyResult,
12+
StoreEncryptedKeyBatchResult,
13+
StoreEncryptedKeyResult,
14+
ImportPrivateKeyResult,
15+
StoredKeyData,
16+
StoredKeyMetadata,
17+
BatchGeneratePrivateKeysResult,
18+
SignTransactionParamsSupportedEvm,
19+
} from '@lit-protocol/wrapped-keys';
20+
import type { SessionSigsMap } from '@lit-protocol/types';
21+
22+
const DEFAULT_NETWORK = wrappedKeysConstants.NETWORK_EVM;
23+
24+
export interface WrappedKeysExecuteJsContext {
25+
litClient: LitClient;
26+
sessionSigs: SessionSigsMap;
27+
}
28+
29+
const ZERO_ADDRESS = '0x0000000000000000000000000000000000000000';
30+
31+
const randomMemo = (prefix: string) =>
32+
`${prefix}-${Date.now()}-${Math.floor(Math.random() * 1e6)}`;
33+
34+
const randomCiphertext = () => randomBytes(48).toString('base64');
35+
36+
const randomHash = (input: string) =>
37+
createHash('sha256').update(input).digest('hex');
38+
39+
const createStorePayload = (memo?: string) => {
40+
const ciphertext = randomCiphertext();
41+
return {
42+
ciphertext,
43+
dataToEncryptHash: randomHash(ciphertext),
44+
publicKey: `0x${randomBytes(33).toString('hex')}`,
45+
keyType: 'K256' as const,
46+
memo: memo ?? randomMemo('store'),
47+
};
48+
};
49+
50+
export async function generateWrappedKey(
51+
context: WrappedKeysExecuteJsContext,
52+
memo = randomMemo('generated')
53+
): Promise<{ memo: string; result: GeneratePrivateKeyResult }> {
54+
const result = await wrappedKeysApi.generatePrivateKey({
55+
pkpSessionSigs: context.sessionSigs,
56+
litClient: context.litClient,
57+
network: DEFAULT_NETWORK,
58+
memo,
59+
});
60+
61+
return { memo, result };
62+
}
63+
64+
export async function listWrappedKeyMetadata(
65+
context: WrappedKeysExecuteJsContext
66+
): Promise<StoredKeyMetadata[]> {
67+
return wrappedKeysApi.listEncryptedKeyMetadata({
68+
pkpSessionSigs: context.sessionSigs,
69+
litClient: context.litClient,
70+
});
71+
}
72+
73+
export async function getWrappedKey(
74+
context: WrappedKeysExecuteJsContext,
75+
id: string
76+
): Promise<StoredKeyData> {
77+
return wrappedKeysApi.getEncryptedKey({
78+
pkpSessionSigs: context.sessionSigs,
79+
litClient: context.litClient,
80+
id,
81+
});
82+
}
83+
84+
export async function exportWrappedKey(
85+
context: WrappedKeysExecuteJsContext,
86+
id: string
87+
): Promise<ExportPrivateKeyResult> {
88+
return wrappedKeysApi.exportPrivateKey({
89+
pkpSessionSigs: context.sessionSigs,
90+
litClient: context.litClient,
91+
id,
92+
network: DEFAULT_NETWORK,
93+
});
94+
}
95+
96+
export async function signMessageWithWrappedKey(
97+
context: WrappedKeysExecuteJsContext,
98+
id: string,
99+
message: string
100+
): Promise<string> {
101+
return wrappedKeysApi.signMessageWithEncryptedKey({
102+
pkpSessionSigs: context.sessionSigs,
103+
litClient: context.litClient,
104+
id,
105+
network: DEFAULT_NETWORK,
106+
messageToSign: message,
107+
});
108+
}
109+
110+
const EVM_NETWORK: SignTransactionParamsSupportedEvm['network'] = 'evm';
111+
112+
export async function signTransactionWithWrappedKey(
113+
context: WrappedKeysExecuteJsContext,
114+
id: string,
115+
unsignedTransaction: SignTransactionParamsSupportedEvm['unsignedTransaction'] = {
116+
toAddress: ZERO_ADDRESS,
117+
value: '0',
118+
chainId: 1,
119+
chain: 'ethereum',
120+
},
121+
broadcast = false
122+
): Promise<string> {
123+
const request: SignTransactionParamsSupportedEvm = {
124+
pkpSessionSigs: context.sessionSigs,
125+
litClient: context.litClient,
126+
id,
127+
network: EVM_NETWORK,
128+
unsignedTransaction,
129+
broadcast,
130+
};
131+
132+
return wrappedKeysApi.signTransactionWithEncryptedKey(request);
133+
}
134+
135+
export async function storeWrappedKey(
136+
context: WrappedKeysExecuteJsContext,
137+
overrides?: Partial<ReturnType<typeof createStorePayload>>
138+
): Promise<{ payload: ReturnType<typeof createStorePayload>; result: StoreEncryptedKeyResult }> {
139+
const payload = { ...createStorePayload(), ...overrides };
140+
141+
const result = await wrappedKeysApi.storeEncryptedKey({
142+
pkpSessionSigs: context.sessionSigs,
143+
litClient: context.litClient,
144+
...payload,
145+
});
146+
147+
return { payload, result };
148+
}
149+
150+
export async function storeWrappedKeyBatch(
151+
context: WrappedKeysExecuteJsContext,
152+
count = 2
153+
): Promise<{
154+
payload: ReturnType<typeof createStorePayload>[];
155+
result: StoreEncryptedKeyBatchResult;
156+
}> {
157+
const payload = Array.from({ length: count }, () => createStorePayload());
158+
159+
const result = await wrappedKeysApi.storeEncryptedKeyBatch({
160+
pkpSessionSigs: context.sessionSigs,
161+
litClient: context.litClient,
162+
keyBatch: payload,
163+
});
164+
165+
return { payload, result };
166+
}
167+
168+
export async function importWrappedKey(
169+
context: WrappedKeysExecuteJsContext,
170+
memo = randomMemo('imported')
171+
): Promise<{ wallet: Wallet; memo: string; result: ImportPrivateKeyResult }> {
172+
const wallet = Wallet.createRandom();
173+
const result = await wrappedKeysApi.importPrivateKey({
174+
pkpSessionSigs: context.sessionSigs,
175+
litClient: context.litClient,
176+
privateKey: wallet.privateKey,
177+
publicKey: wallet.publicKey,
178+
keyType: 'K256',
179+
memo,
180+
});
181+
182+
return { wallet, memo, result };
183+
}
184+
185+
export async function batchGenerateWrappedKeys(
186+
context: WrappedKeysExecuteJsContext,
187+
count = 2
188+
): Promise<BatchGeneratePrivateKeysResult> {
189+
const actions = Array.from({ length: count }, (_, idx) => {
190+
const memo = randomMemo(`batch-${idx}`);
191+
return {
192+
network: DEFAULT_NETWORK,
193+
generateKeyParams: { memo },
194+
...(idx === 0
195+
? { signMessageParams: { messageToSign: `batch-${idx}-message` } }
196+
: {}),
197+
};
198+
});
199+
200+
return wrappedKeysApi.batchGeneratePrivateKeys({
201+
pkpSessionSigs: context.sessionSigs,
202+
litClient: context.litClient,
203+
actions,
204+
});
205+
}

0 commit comments

Comments
 (0)