/**
* Unit tests for Code Mode Tool (pg_execute_code)
*
* Tests the tool handler, initialization, and cleanup functions.
*/
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import {
createExecuteCodeTool,
getCodeModeTools,
cleanupCodeMode
} from '../codemode/index.js';
import type { PostgresAdapter } from '../../PostgresAdapter.js';
// Mock the dependencies
vi.mock('../../../codemode/sandbox-factory.js', () => ({
createSandboxPool: vi.fn(() => ({
initialize: vi.fn(),
execute: vi.fn().mockResolvedValue({
success: true,
result: { test: 'result' },
metrics: { wallTimeMs: 10, cpuTimeMs: 8, memoryUsedMb: 1 }
}),
dispose: vi.fn(),
getStats: vi.fn().mockReturnValue({ available: 5, inUse: 0, max: 5 })
}))
}));
vi.mock('../../../codemode/security.js', () => ({
CodeModeSecurityManager: vi.fn().mockImplementation(() => ({
validateCode: vi.fn().mockReturnValue({ valid: true, errors: [] }),
checkRateLimit: vi.fn().mockReturnValue(true),
sanitizeResult: vi.fn((result: unknown) => result),
createExecutionRecord: vi.fn().mockReturnValue({}),
auditLog: vi.fn()
}))
}));
vi.mock('../../../codemode/api.js', () => ({
createPgApi: vi.fn().mockReturnValue({
createSandboxBindings: vi.fn().mockReturnValue({ core: { listTables: () => [] } })
})
}));
vi.mock('../../../utils/icons.js', () => ({
getToolIcons: vi.fn().mockReturnValue([])
}));
vi.mock('../../../utils/logger.js', () => ({
logger: {
debug: vi.fn(),
info: vi.fn(),
warn: vi.fn(),
error: vi.fn()
}
}));
/**
* Create a mock PostgresAdapter for testing
*/
function createMockAdapter(): Partial<PostgresAdapter> {
return {
type: 'postgresql' as const,
name: 'PostgreSQL Adapter',
version: '0.1.0',
executeQuery: vi.fn().mockResolvedValue({ rows: [], rowsAffected: 0, executionTimeMs: 1 }),
executeReadQuery: vi.fn().mockResolvedValue({ rows: [], rowsAffected: 0, executionTimeMs: 1 }),
executeWriteQuery: vi.fn().mockResolvedValue({ rows: [], rowsAffected: 0, executionTimeMs: 1 }),
getToolDefinitions: vi.fn().mockReturnValue([])
};
}
describe('Code Mode Tool', () => {
let mockAdapter: Partial<PostgresAdapter>;
beforeEach(() => {
mockAdapter = createMockAdapter();
// Clear any existing singleton state
cleanupCodeMode();
vi.clearAllMocks();
});
afterEach(() => {
cleanupCodeMode();
});
describe('createExecuteCodeTool', () => {
it('should return a valid tool definition', () => {
const tool = createExecuteCodeTool(mockAdapter as PostgresAdapter);
expect(tool).toBeDefined();
expect(tool.name).toBe('pg_execute_code');
});
it('should have correct tool group', () => {
const tool = createExecuteCodeTool(mockAdapter as PostgresAdapter);
expect(tool.group).toBe('codemode');
});
it('should have correct tags', () => {
const tool = createExecuteCodeTool(mockAdapter as PostgresAdapter);
expect(tool.tags).toContain('code');
expect(tool.tags).toContain('execute');
expect(tool.tags).toContain('sandbox');
});
it('should require admin scope', () => {
const tool = createExecuteCodeTool(mockAdapter as PostgresAdapter);
expect(tool.requiredScopes).toContain('admin');
});
it('should have destructive annotation', () => {
const tool = createExecuteCodeTool(mockAdapter as PostgresAdapter);
expect(tool.annotations?.destructiveHint).toBe(true);
});
it('should have readOnlyHint as false', () => {
const tool = createExecuteCodeTool(mockAdapter as PostgresAdapter);
expect(tool.annotations?.readOnlyHint).toBe(false);
});
it('should have description mentioning pg.* API', () => {
const tool = createExecuteCodeTool(mockAdapter as PostgresAdapter);
expect(tool.description).toContain('pg.*');
});
it('should include input schema', () => {
const tool = createExecuteCodeTool(mockAdapter as PostgresAdapter);
expect(tool.inputSchema).toBeDefined();
});
});
describe('getCodeModeTools', () => {
it('should return array containing pg_execute_code', () => {
const tools = getCodeModeTools(mockAdapter as PostgresAdapter);
expect(tools.length).toBeGreaterThan(0);
expect(tools.some(t => t.name === 'pg_execute_code')).toBe(true);
});
it('should return exactly one tool', () => {
const tools = getCodeModeTools(mockAdapter as PostgresAdapter);
expect(tools.length).toBe(1);
});
});
describe('cleanupCodeMode', () => {
it('should dispose pool when called', () => {
// Initialize by creating a tool and calling handler
const tool = createExecuteCodeTool(mockAdapter as PostgresAdapter);
expect(tool).toBeDefined();
// Cleanup should not throw
expect(() => cleanupCodeMode()).not.toThrow();
});
it('should be idempotent', () => {
cleanupCodeMode();
cleanupCodeMode();
// No error means success
});
it('should allow re-initialization after cleanup', () => {
const tool1 = createExecuteCodeTool(mockAdapter as PostgresAdapter);
cleanupCodeMode();
const tool2 = createExecuteCodeTool(mockAdapter as PostgresAdapter);
expect(tool1.name).toBe(tool2.name);
});
});
describe('tool handler', () => {
it('should have a handler function', () => {
const tool = createExecuteCodeTool(mockAdapter as PostgresAdapter);
expect(typeof tool.handler).toBe('function');
});
it('should execute valid code successfully', async () => {
const tool = createExecuteCodeTool(mockAdapter as PostgresAdapter);
const result = await tool.handler({ code: 'return 1 + 1' }, { timestamp: new Date(), requestId: 'test' });
expect(result).toHaveProperty('success');
});
});
});
// Note: getIsolationMode is tested implicitly through sandbox-factory tests
