File size: 3,593 Bytes
84aa3bf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
/**
 * @license
 * Copyright 2026 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import { vi } from 'vitest';
import type {
  ShellExecutionHandle,
  ShellExecutionResult,
  ShellOutputEvent,
  ShellExecutionConfig,
} from '@google/gemini-cli-core';

export interface MockShellCommand {
  command: string | RegExp;
  result: Partial<ShellExecutionResult>;
  events?: ShellOutputEvent[];
}

type ShellExecutionServiceExecute = (
  commandToExecute: string,
  cwd: string,
  onOutputEvent: (event: ShellOutputEvent) => void,
  abortSignal: AbortSignal,
  shouldUseNodePty: boolean,
  shellExecutionConfig: ShellExecutionConfig,
) => Promise<ShellExecutionHandle>;

export class MockShellExecutionService {
  private static mockCommands: MockShellCommand[] = [];
  private static originalExecute: ShellExecutionServiceExecute | undefined;
  private static passthroughEnabled = false;

  /**
   * Registers the original implementation to allow falling back to real shell execution.
   */
  static setOriginalImplementation(
    implementation: ShellExecutionServiceExecute,
  ) {
    this.originalExecute = implementation;
  }

  /**
   * Enables or disables passthrough to the real implementation when no mock matches.
   */
  static setPassthrough(enabled: boolean) {
    this.passthroughEnabled = enabled;
  }

  static setMockCommands(commands: MockShellCommand[]) {
    this.mockCommands = commands;
  }

  static reset() {
    this.mockCommands = [];
    this.passthroughEnabled = false;
    this.writeToPty.mockClear();
    this.kill.mockClear();
    this.background.mockClear();
    this.resizePty.mockClear();
    this.scrollPty.mockClear();
  }

  static async execute(
    commandToExecute: string,
    cwd: string,
    onOutputEvent: (event: ShellOutputEvent) => void,
    abortSignal: AbortSignal,
    shouldUseNodePty: boolean,
    shellExecutionConfig: ShellExecutionConfig,
  ): Promise<ShellExecutionHandle> {
    const mock = this.mockCommands.find((m) =>
      typeof m.command === 'string'
        ? m.command === commandToExecute
        : m.command.test(commandToExecute),
    );

    const pid = Math.floor(Math.random() * 10000);

    if (mock) {
      if (mock.events) {
        for (const event of mock.events) {
          onOutputEvent(event);
        }
      }

      const result: ShellExecutionResult = {
        rawOutput: Buffer.from(mock.result.output || ''),
        output: mock.result.output || '',
        exitCode: mock.result.exitCode ?? 0,
        signal: mock.result.signal ?? null,
        error: mock.result.error ?? null,
        aborted: false,
        pid,
        executionMethod: 'none',
        ...mock.result,
      };

      return {
        pid,
        result: Promise.resolve(result),
      };
    }

    if (this.passthroughEnabled && this.originalExecute) {
      return this.originalExecute(
        commandToExecute,
        cwd,
        onOutputEvent,
        abortSignal,
        shouldUseNodePty,
        shellExecutionConfig,
      );
    }

    return {
      pid,
      result: Promise.resolve({
        rawOutput: Buffer.from(''),
        output: `Command not found: ${commandToExecute}`,
        exitCode: 127,
        signal: null,
        error: null,
        aborted: false,
        pid,
        executionMethod: 'none',
      }),
    };
  }

  static writeToPty = vi.fn();
  static isPtyActive = vi.fn(() => false);
  static onExit = vi.fn(() => () => {});
  static kill = vi.fn();
  static background = vi.fn();
  static subscribe = vi.fn(() => () => {});
  static resizePty = vi.fn();
  static scrollPty = vi.fn();
}