Hugging Face's logo

SaylorTwift
/
gemini-cli

agent
coding-agent
cli
llm
tool-use
Model card Files
xet
 Community
gemini-cli / packages /core /src /test-utils /mock-message-bus.ts
SaylorTwift's picture
SaylorTwift HF Staff
Add files using upload-large-folder tool
52efc7b verified
Raw History Blame Contribute Delete
3.24 kB
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { vi } from 'vitest';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import { MessageBusType, type Message } from '../confirmation-bus/types.js';
/**
* Mock MessageBus for testing hook execution through MessageBus
*/
export class MockMessageBus {
private subscriptions = new Map<
MessageBusType,
Set<(message: Message) => void>
>();
publishedMessages: Message[] = [];
defaultToolDecision: 'allow' | 'deny' | 'ask_user' = 'allow';
/**
* Mock publish method that captures messages and simulates responses
*/
publish = vi.fn(async (message: Message) => {
this.publishedMessages.push(message);
// Handle tool confirmation requests
if (message.type === MessageBusType.TOOL_CONFIRMATION_REQUEST) {
if (this.defaultToolDecision === 'allow') {
this.emit(MessageBusType.TOOL_CONFIRMATION_RESPONSE, {
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
correlationId: message.correlationId,
confirmed: true,
});
} else if (this.defaultToolDecision === 'deny') {
this.emit(MessageBusType.TOOL_CONFIRMATION_RESPONSE, {
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
correlationId: message.correlationId,
confirmed: false,
});
} else {
// ask_user
this.emit(MessageBusType.TOOL_CONFIRMATION_RESPONSE, {
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
correlationId: message.correlationId,
confirmed: false,
requiresUserConfirmation: true,
});
}
}
// Emit the message to subscribers (mimicking real MessageBus behavior)
this.emit(message.type, message);
});
/**
* Mock subscribe method that stores listeners
*/
subscribe = vi.fn(
<T extends Message>(type: T['type'], listener: (message: T) => void) => {
if (!this.subscriptions.has(type)) {
this.subscriptions.set(type, new Set());
}
this.subscriptions.get(type)!.add(listener as (message: Message) => void);
},
);
/**
* Mock unsubscribe method
*/
unsubscribe = vi.fn(
<T extends Message>(type: T['type'], listener: (message: T) => void) => {
const listeners = this.subscriptions.get(type);
if (listeners) {
listeners.delete(listener as (message: Message) => void);
}
},
);
/**
* Emit a message to subscribers (for testing)
*/
private emit(type: MessageBusType, message: Message) {
const listeners = this.subscriptions.get(type);
if (listeners) {
listeners.forEach((listener) => listener(message));
}
}
/**
* Clear all captured messages (for test isolation)
*/
clear() {
this.publishedMessages = [];
this.subscriptions.clear();
}
}
/**
* Create a mock MessageBus for testing
*/
export function createMockMessageBus(): MessageBus {
return new MockMessageBus() as unknown as MessageBus;
}
/**
* Get the MockMessageBus instance from a mocked MessageBus
*/
export function getMockMessageBusInstance(
messageBus: MessageBus,
): MockMessageBus {
return messageBus as unknown as MockMessageBus;
}