Skip to content

Commit

Permalink
test: plugin-tee - adjusting project structure and new tests (#2508)
Browse files Browse the repository at this point in the history
* test: plugin-tee - adjusting project strucutre and new tests

* test: plugin-tee - ficed test structure

---------

Co-authored-by: Sayo <[email protected]>
  • Loading branch information
ai16z-demirix and wtfsayo authored Jan 19, 2025
1 parent ace85ae commit 319840c
Show file tree
Hide file tree
Showing 4 changed files with 130 additions and 35 deletions.
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { DeriveKeyProvider } from '../providers/deriveKeyProvider';
import { DeriveKeyProvider } from '../src/providers/deriveKeyProvider';
import { TappdClient } from '@phala/dstack-sdk';
import { TEEMode } from '../types/tee';
import { TEEMode } from '../src/types/tee';

// Mock dependencies
vi.mock('@phala/dstack-sdk', () => ({
Expand Down Expand Up @@ -101,33 +101,11 @@ describe('DeriveKeyProvider', () => {
it('should derive Ed25519 keypair successfully', async () => {
const path = 'test-path';
const subject = 'test-subject';
const agentId = 'test-agent';
const result = await provider.deriveEd25519Keypair(path, subject);

const result = await provider.deriveEd25519Keypair(path, subject, agentId);

expect(result).toHaveProperty('keypair');
expect(result).toHaveProperty('attestation');
expect(result.keypair.publicKey.toBase58()).toBe('mock-solana-public-key');
});
});

describe('deriveEcdsaKeypair', () => {
let provider: DeriveKeyProvider;

beforeEach(() => {
provider = new DeriveKeyProvider(TEEMode.LOCAL);
});

it('should derive ECDSA keypair successfully', async () => {
const path = 'test-path';
const subject = 'test-subject';
const agentId = 'test-agent';

const result = await provider.deriveEcdsaKeypair(path, subject, agentId);

expect(result).toHaveProperty('keypair');
expect(result).toHaveProperty('attestation');
expect(result.keypair.address).toBe('mock-evm-address');
const client = TappdClient.mock.results[0].value;
expect(client.deriveKey).toHaveBeenCalledWith(path, subject);
expect(result.keypair.publicKey.toBase58()).toEqual('mock-solana-public-key');
});
});
});
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { RemoteAttestationProvider } from '../providers/remoteAttestationProvider';
import { RemoteAttestationProvider } from '../src/providers/remoteAttestationProvider';
import { TappdClient } from '@phala/dstack-sdk';
import { TEEMode } from '../types/tee';
import { TEEMode } from '../src/types/tee';

// Mock TappdClient
vi.mock('@phala/dstack-sdk', () => ({
Expand Down
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { remoteAttestationAction } from '../actions/remoteAttestation';
import { RemoteAttestationProvider } from '../providers/remoteAttestationProvider';
import { remoteAttestationAction } from '../src/actions/remoteAttestation';
import { RemoteAttestationProvider } from '../src/providers/remoteAttestationProvider';

// Mock dependencies
vi.mock('../providers/remoteAttestationProvider');
vi.mock('../src/providers/remoteAttestationProvider');
vi.mock('undici', () => ({
fetch: vi.fn().mockResolvedValue({
json: () => Promise.resolve({ checksum: 'mock-checksum' })
Expand Down Expand Up @@ -76,7 +76,6 @@ describe('remoteAttestationAction', () => {

expect(result).toBe(false);
});

});

describe('validate', () => {
Expand All @@ -94,10 +93,11 @@ describe('remoteAttestationAction', () => {
const [userMessage, agentMessage] = remoteAttestationAction.examples[0];
expect(userMessage.user).toBe('{{user1}}');
expect(userMessage.content.text).toBe('If you are running in a TEE, generate a remote attestation');
expect(userMessage.content.action).toBe('REMOTE_ATTESTATION');

expect(agentMessage.user).toBe('{{user2}}');
expect(agentMessage.content.text).toBe('Of course, one second...');
expect(agentMessage.content.action).toBe('REMOTE_ATTESTATION');
expect(agentMessage.content.action).toBeUndefined();
});
});
});
117 changes: 117 additions & 0 deletions packages/plugin-tee/__tests__/timeout.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,117 @@
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { RemoteAttestationProvider } from '../src/providers/remoteAttestationProvider';
import { DeriveKeyProvider } from '../src/providers/deriveKeyProvider';
import { TEEMode } from '../src/types/tee';
import { TappdClient } from '@phala/dstack-sdk';

// Mock TappdClient
vi.mock('@phala/dstack-sdk', () => ({
TappdClient: vi.fn().mockImplementation(() => ({
tdxQuote: vi.fn(),
deriveKey: vi.fn()
}))
}));

describe('TEE Provider Timeout Tests', () => {
beforeEach(() => {
vi.clearAllMocks();
});

describe('RemoteAttestationProvider', () => {
it('should handle API timeout during attestation generation', async () => {
const mockTdxQuote = vi.fn()
.mockRejectedValueOnce(new Error('Request timed out'));

vi.mocked(TappdClient).mockImplementation(() => ({
tdxQuote: mockTdxQuote,
deriveKey: vi.fn()
}));

const provider = new RemoteAttestationProvider(TEEMode.LOCAL);
await expect(() => provider.generateAttestation('test-data'))
.rejects
.toThrow('Failed to generate TDX Quote: Request timed out');

// Verify the call was made once
expect(mockTdxQuote).toHaveBeenCalledTimes(1);
expect(mockTdxQuote).toHaveBeenCalledWith('test-data', undefined);
});

it('should handle network errors during attestation generation', async () => {
const mockTdxQuote = vi.fn()
.mockRejectedValueOnce(new Error('Network error'));

vi.mocked(TappdClient).mockImplementation(() => ({
tdxQuote: mockTdxQuote,
deriveKey: vi.fn()
}));

const provider = new RemoteAttestationProvider(TEEMode.LOCAL);
await expect(() => provider.generateAttestation('test-data'))
.rejects
.toThrow('Failed to generate TDX Quote: Network error');

expect(mockTdxQuote).toHaveBeenCalledTimes(1);
});

it('should handle successful attestation generation', async () => {
const mockQuote = {
quote: 'test-quote',
replayRtmrs: () => ['rtmr0', 'rtmr1', 'rtmr2', 'rtmr3']
};

const mockTdxQuote = vi.fn().mockResolvedValueOnce(mockQuote);

vi.mocked(TappdClient).mockImplementation(() => ({
tdxQuote: mockTdxQuote,
deriveKey: vi.fn()
}));

const provider = new RemoteAttestationProvider(TEEMode.LOCAL);
const result = await provider.generateAttestation('test-data');

expect(mockTdxQuote).toHaveBeenCalledTimes(1);
expect(result).toEqual({
quote: 'test-quote',
timestamp: expect.any(Number)
});
});
});

describe('DeriveKeyProvider', () => {
it('should handle API timeout during key derivation', async () => {
const mockDeriveKey = vi.fn()
.mockRejectedValueOnce(new Error('Request timed out'));

vi.mocked(TappdClient).mockImplementation(() => ({
tdxQuote: vi.fn(),
deriveKey: mockDeriveKey
}));

const provider = new DeriveKeyProvider(TEEMode.LOCAL);
await expect(() => provider.rawDeriveKey('test-path', 'test-subject'))
.rejects
.toThrow('Request timed out');

expect(mockDeriveKey).toHaveBeenCalledTimes(1);
expect(mockDeriveKey).toHaveBeenCalledWith('test-path', 'test-subject');
});

it('should handle API timeout during Ed25519 key derivation', async () => {
const mockDeriveKey = vi.fn()
.mockRejectedValueOnce(new Error('Request timed out'));

vi.mocked(TappdClient).mockImplementation(() => ({
tdxQuote: vi.fn(),
deriveKey: mockDeriveKey
}));

const provider = new DeriveKeyProvider(TEEMode.LOCAL);
await expect(() => provider.deriveEd25519Keypair('test-path', 'test-subject'))
.rejects
.toThrow('Request timed out');

expect(mockDeriveKey).toHaveBeenCalledTimes(1);
});
});
});

0 comments on commit 319840c

Please sign in to comment.