Lupul-Augmentat/tests/tools/SystemCommandTool.test.ts

import { SystemCommandTool } from '../../src/tools/builtin/SystemCommandTool';
import { ToolContext } from '../../src/tools/base/ToolHandler';
import { spawn } from 'child_process';
import { EventEmitter } from 'events';

// Mock child_process
jest.mock('child_process');

class MockChildProcess extends EventEmitter {
  stdout = new EventEmitter();
  stderr = new EventEmitter();
  stdin = {
    write: jest.fn(),
    end: jest.fn(),
  };
  kill = jest.fn();
}

describe('SystemCommandTool', () => {
  jest.setTimeout(10000); // Increase timeout for these tests
  let tool: SystemCommandTool;
  let mockSpawn: jest.MockedFunction<typeof spawn>;

  beforeEach(() => {
    tool = new SystemCommandTool();
    mockSpawn = spawn as jest.MockedFunction<typeof spawn>;
    jest.clearAllMocks();
  });

  const createContext = (permissions: string[] = ['system:exec']): ToolContext => ({
    requestId: 'test-request',
    permissions,
  });

  describe('execute', () => {
    it('should execute allowed commands', async () => {
      const mockProcess = new MockChildProcess();
      mockSpawn.mockReturnValue(mockProcess as any);

      const resultPromise = tool.execute(
        { command: 'ls', args: ['-la'] },
        createContext(),
      );

      // Simulate command output
      setImmediate(() => {
        mockProcess.stdout.emit('data', Buffer.from('file1.txt\nfile2.txt'));
        mockProcess.emit('close', 0);
      });

      const result = await resultPromise;

      expect(result).toMatchObject({
        stdout: 'file1.txt\nfile2.txt',
        stderr: '',
        exitCode: 0,
      });

      expect(mockSpawn).toHaveBeenCalledWith('ls', ['-la'], expect.any(Object));
    });

    it('should handle command with environment variables', async () => {
      const mockProcess = new MockChildProcess();
      mockSpawn.mockReturnValue(mockProcess as any);

      const resultPromise = tool.execute(
        { 
          command: 'echo', 
          args: ['test'],
          env: { TEST_VAR: 'test-value' },
        },
        createContext(),
      );

      setImmediate(() => {
        mockProcess.stdout.emit('data', Buffer.from('test'));
        mockProcess.emit('close', 0);
      });

      await resultPromise;

      expect(mockSpawn).toHaveBeenCalledWith(
        'echo',
        ['test'],
        expect.objectContaining({
          env: expect.objectContaining({
            TEST_VAR: 'test-value',
          }),
        }),
      );
    });

    it('should reject non-whitelisted commands', async () => {
      await expect(
        tool.execute(
          { command: 'rm', args: ['-rf', '/'] },
          createContext(),
        ),
      ).rejects.toThrow('Command not allowed: rm');
    });

    it('should prevent command injection', async () => {
      await expect(
        tool.execute(
          { command: 'ls', args: ['; rm -rf /'] },
          createContext(),
        ),
      ).rejects.toThrow('Shell characters not allowed in');
    });

    it('should handle command timeout', async () => {
      const mockProcess = new MockChildProcess();
      mockSpawn.mockReturnValue(mockProcess as any);

      const resultPromise = tool.execute(
        { command: 'ls', timeout: 100 },
        createContext(),
      );

      // Don't emit close event to simulate timeout
      await expect(resultPromise).rejects.toThrow('Command timed out after 100ms');
    });

    it('should handle command failure', async () => {
      const mockProcess = new MockChildProcess();
      mockSpawn.mockReturnValue(mockProcess as any);

      const resultPromise = tool.execute(
        { command: 'ls' },
        createContext(),
      );

      setImmediate(() => {
        mockProcess.stderr.emit('data', Buffer.from('Command not found'));
        mockProcess.emit('close', 127);
      });

      const result = await resultPromise;

      expect(result).toEqual({
        stdout: '',
        stderr: 'Command not found',
        exitCode: 127,
        duration: expect.any(Number),
      });
    });

    it('should require system:exec permission', async () => {
      await expect(
        tool.execute(
          { command: 'ls' },
          createContext([]),
        ),
      ).rejects.toThrow('Permission denied: system:exec required');
    });

    it('should respect working directory', async () => {
      const mockProcess = new MockChildProcess();
      mockSpawn.mockReturnValue(mockProcess as any);

      const resultPromise = tool.execute(
        { command: 'ls', cwd: '/tmp' },
        createContext(),
      );

      setImmediate(() => {
        mockProcess.emit('close', 0);
      });

      await resultPromise;

      expect(mockSpawn).toHaveBeenCalledWith(
        'ls',
        [],
        expect.objectContaining({
          cwd: '/tmp',
        }),
      );
    });

    it('should handle stdin input', async () => {
      const mockProcess = new MockChildProcess();
      mockSpawn.mockReturnValue(mockProcess as any);

      const resultPromise = tool.execute(
        { command: 'cat', stdin: 'Hello, World!' },
        createContext(),
      );

      setImmediate(() => {
        mockProcess.stdout.emit('data', Buffer.from('Hello, World!'));
        mockProcess.emit('close', 0);
      });

      await resultPromise;

      expect(mockProcess.stdin.write).toHaveBeenCalledWith('Hello, World!');
      expect(mockProcess.stdin.end).toHaveBeenCalled();
    });
  });
});