1+ import { describe , it , expect , vi , beforeEach } from 'vitest' ;
2+ import { DeriveKeyProvider } from '../providers/deriveKeyProvider' ;
3+ import { TappdClient } from '@phala/dstack-sdk' ;
4+ import { TEEMode } from '../types/tee' ;
5+ import { Keypair } from '@solana/web3.js' ;
6+ import { privateKeyToAccount } from 'viem/accounts' ;
7+
8+ // Mock dependencies
9+ vi . mock ( '@phala/dstack-sdk' , ( ) => ( {
10+ TappdClient : vi . fn ( ) . mockImplementation ( ( ) => ( {
11+ deriveKey : vi . fn ( ) . mockResolvedValue ( {
12+ asUint8Array : ( ) => new Uint8Array ( [ 1 , 2 , 3 , 4 , 5 ] )
13+ } ) ,
14+ tdxQuote : vi . fn ( ) . mockResolvedValue ( {
15+ quote : 'mock-quote-data' ,
16+ replayRtmrs : ( ) => [ 'rtmr0' , 'rtmr1' , 'rtmr2' , 'rtmr3' ]
17+ } ) ,
18+ rawDeriveKey : vi . fn ( )
19+ } ) )
20+ } ) ) ;
21+
22+ vi . mock ( '@solana/web3.js' , ( ) => ( {
23+ Keypair : {
24+ fromSeed : vi . fn ( ) . mockReturnValue ( {
25+ publicKey : {
26+ toBase58 : ( ) => 'mock-solana-public-key'
27+ }
28+ } )
29+ }
30+ } ) ) ;
31+
32+ vi . mock ( 'viem/accounts' , ( ) => ( {
33+ privateKeyToAccount : vi . fn ( ) . mockReturnValue ( {
34+ address : 'mock-evm-address'
35+ } )
36+ } ) ) ;
37+
38+ describe ( 'DeriveKeyProvider' , ( ) => {
39+ beforeEach ( ( ) => {
40+ vi . clearAllMocks ( ) ;
41+ } ) ;
42+
43+ describe ( 'constructor' , ( ) => {
44+ it ( 'should initialize with LOCAL mode' , ( ) => {
45+ const _provider = new DeriveKeyProvider ( TEEMode . LOCAL ) ;
46+ expect ( TappdClient ) . toHaveBeenCalledWith ( 'http://localhost:8090' ) ;
47+ } ) ;
48+
49+ it ( 'should initialize with DOCKER mode' , ( ) => {
50+ const _provider = new DeriveKeyProvider ( TEEMode . DOCKER ) ;
51+ expect ( TappdClient ) . toHaveBeenCalledWith ( 'http://host.docker.internal:8090' ) ;
52+ } ) ;
53+
54+ it ( 'should initialize with PRODUCTION mode' , ( ) => {
55+ const _provider = new DeriveKeyProvider ( TEEMode . PRODUCTION ) ;
56+ expect ( TappdClient ) . toHaveBeenCalledWith ( ) ;
57+ } ) ;
58+
59+ it ( 'should throw error for invalid mode' , ( ) => {
60+ expect ( ( ) => new DeriveKeyProvider ( 'INVALID_MODE' ) ) . toThrow ( 'Invalid TEE_MODE' ) ;
61+ } ) ;
62+ } ) ;
63+
64+ describe ( 'rawDeriveKey' , ( ) => {
65+ let _provider : DeriveKeyProvider ;
66+
67+ beforeEach ( ( ) => {
68+ _provider = new DeriveKeyProvider ( TEEMode . LOCAL ) ;
69+ } ) ;
70+
71+ it ( 'should derive raw key successfully' , async ( ) => {
72+ const path = 'test-path' ;
73+ const subject = 'test-subject' ;
74+ const result = await _provider . rawDeriveKey ( path , subject ) ;
75+
76+ const client = TappdClient . mock . results [ 0 ] . value ;
77+ expect ( client . deriveKey ) . toHaveBeenCalledWith ( path , subject ) ;
78+ expect ( result . asUint8Array ( ) ) . toEqual ( new Uint8Array ( [ 1 , 2 , 3 , 4 , 5 ] ) ) ;
79+ } ) ;
80+
81+ it ( 'should handle errors during raw key derivation' , async ( ) => {
82+ const mockError = new Error ( 'Key derivation failed' ) ;
83+ vi . mocked ( TappdClient ) . mockImplementationOnce ( ( ) => {
84+ const instance = new TappdClient ( ) ;
85+ instance . deriveKey = vi . fn ( ) . mockRejectedValueOnce ( mockError ) ;
86+ instance . tdxQuote = vi . fn ( ) ;
87+ instance . rawDeriveKey = vi . fn ( ) ;
88+ return instance ;
89+ } ) ;
90+
91+ const provider = new DeriveKeyProvider ( TEEMode . LOCAL ) ;
92+ await expect ( provider . rawDeriveKey ( 'path' , 'subject' ) ) . rejects . toThrow ( mockError ) ;
93+ } ) ;
94+ } ) ;
95+
96+ describe ( 'deriveEd25519Keypair' , ( ) => {
97+ let provider : DeriveKeyProvider ;
98+
99+ beforeEach ( ( ) => {
100+ provider = new DeriveKeyProvider ( TEEMode . LOCAL ) ;
101+ } ) ;
102+
103+ it ( 'should derive Ed25519 keypair successfully' , async ( ) => {
104+ const path = 'test-path' ;
105+ const subject = 'test-subject' ;
106+ const agentId = 'test-agent' ;
107+
108+ const result = await provider . deriveEd25519Keypair ( path , subject , agentId ) ;
109+
110+ expect ( result ) . toHaveProperty ( 'keypair' ) ;
111+ expect ( result ) . toHaveProperty ( 'attestation' ) ;
112+ expect ( result . keypair . publicKey . toBase58 ( ) ) . toBe ( 'mock-solana-public-key' ) ;
113+ } ) ;
114+ } ) ;
115+
116+ describe ( 'deriveEcdsaKeypair' , ( ) => {
117+ let provider : DeriveKeyProvider ;
118+
119+ beforeEach ( ( ) => {
120+ provider = new DeriveKeyProvider ( TEEMode . LOCAL ) ;
121+ } ) ;
122+
123+ it ( 'should derive ECDSA keypair successfully' , async ( ) => {
124+ const path = 'test-path' ;
125+ const subject = 'test-subject' ;
126+ const agentId = 'test-agent' ;
127+
128+ const result = await provider . deriveEcdsaKeypair ( path , subject , agentId ) ;
129+
130+ expect ( result ) . toHaveProperty ( 'keypair' ) ;
131+ expect ( result ) . toHaveProperty ( 'attestation' ) ;
132+ expect ( result . keypair . address ) . toBe ( 'mock-evm-address' ) ;
133+ } ) ;
134+ } ) ;
135+ } ) ;
0 commit comments