message-handler.test.ts•18.4 kB
import { MessageHandler } from "../message-handler";
import { WebsocketClient } from "../client";
import type { ServerMessageRequest } from "@browser-control-mcp/common";
import { ExtensionConfig } from "../extension-config";
// Mock the WebsocketClient
jest.mock("../client", () => {
return {
WebsocketClient: jest.fn().mockImplementation(() => {
return {
sendResourceToServer: jest.fn().mockResolvedValue(undefined),
sendErrorToServer: jest.fn().mockResolvedValue(undefined),
};
}),
};
});
describe("MessageHandler", () => {
let messageHandler: MessageHandler;
let mockClient: jest.Mocked<WebsocketClient>;
beforeEach(() => {
// Clear all mocks before each test
jest.clearAllMocks();
// Create a new instance of WebsocketClient and MessageHandler
mockClient = new WebsocketClient(
8080,
"test-secret"
) as jest.Mocked<WebsocketClient>;
messageHandler = new MessageHandler(mockClient);
// Mock browser.storage.local.get to return default config
const defaultConfig: ExtensionConfig = {
secret: "test-secret",
toolSettings: {
"open-browser-tab": true,
"close-browser-tabs": true,
"get-list-of-open-tabs": true,
"get-recent-browser-history": true,
"get-tab-web-content": true,
"reorder-browser-tabs": true,
"find-highlight-in-browser-tab": true,
},
domainDenyList: [],
ports: [8089],
auditLog: [],
};
(browser.storage.local.get as jest.Mock).mockResolvedValue({
config: defaultConfig,
});
});
describe("handleDecodedMessage", () => {
it("should throw an error if command is not allowed", async () => {
// Arrange
const configWithDisabledOpenTab: ExtensionConfig = {
secret: "test-secret",
toolSettings: {
"open-browser-tab": false, // Disable open-tab command
"close-browser-tabs": true,
"get-list-of-open-tabs": true,
"get-recent-browser-history": true,
"get-tab-web-content": true,
"reorder-browser-tabs": true,
"find-highlight-in-browser-tab": true,
},
domainDenyList: [],
ports: [8089],
auditLog: [],
};
(browser.storage.local.get as jest.Mock).mockResolvedValue({
config: configWithDisabledOpenTab,
});
const request: ServerMessageRequest = {
cmd: "open-tab",
url: "https://example.com",
correlationId: "test-correlation-id",
};
// Act & Assert
await expect(
messageHandler.handleDecodedMessage(request)
).rejects.toThrow("Command 'open-tab' is disabled in extension settings");
});
describe("open-tab command", () => {
it("should open a new tab and send the tab ID to the server", async () => {
// Arrange
const request: ServerMessageRequest = {
cmd: "open-tab",
url: "https://example.com",
correlationId: "test-correlation-id",
};
const mockTab = { id: 123 };
(browser.tabs.create as jest.Mock).mockResolvedValue(mockTab);
// Act
await messageHandler.handleDecodedMessage(request);
// Assert
expect(browser.tabs.create).toHaveBeenCalledWith({
url: "https://example.com",
});
expect(mockClient.sendResourceToServer).toHaveBeenCalledWith({
resource: "opened-tab-id",
correlationId: "test-correlation-id",
tabId: 123,
});
});
it("should throw an error if URL does not start with https://", async () => {
// Arrange
const request: ServerMessageRequest = {
cmd: "open-tab",
url: "http://example.com",
correlationId: "test-correlation-id",
};
// Act & Assert
await expect(
messageHandler.handleDecodedMessage(request)
).rejects.toThrow("Invalid URL");
expect(browser.tabs.create).not.toHaveBeenCalled();
});
it("should throw an error if domain is in deny list", async () => {
// Arrange
const configWithDenyList: ExtensionConfig = {
secret: "test-secret",
toolSettings: {
"open-browser-tab": true,
"close-browser-tabs": true,
"get-list-of-open-tabs": true,
"get-recent-browser-history": true,
"get-tab-web-content": true,
"reorder-browser-tabs": true,
"find-highlight-in-browser-tab": true,
},
domainDenyList: ["example.com", "another.com"],
ports: [8089],
auditLog: [],
};
(browser.storage.local.get as jest.Mock).mockResolvedValue({
config: configWithDenyList,
});
const request: ServerMessageRequest = {
cmd: "open-tab",
url: "https://example.com",
correlationId: "test-correlation-id",
};
// Act & Assert
await expect(
messageHandler.handleDecodedMessage(request)
).rejects.toThrow("Domain in user defined deny list");
expect(browser.tabs.create).not.toHaveBeenCalled();
});
it("should open a new tab in the domain is not in the deny list", async () => {
// Arrange
const configWithDenyList: ExtensionConfig = {
secret: "test-secret",
toolSettings: {
"open-browser-tab": true,
"close-browser-tabs": true,
"get-list-of-open-tabs": true,
"get-recent-browser-history": true,
"get-tab-web-content": true,
"reorder-browser-tabs": true,
"find-highlight-in-browser-tab": true,
},
domainDenyList: ["example.com", "another.com"],
ports: [8089],
auditLog: [],
};
(browser.storage.local.get as jest.Mock).mockResolvedValue({
config: configWithDenyList,
});
const request: ServerMessageRequest = {
cmd: "open-tab",
url: "https://allowed.com",
correlationId: "test-correlation-id",
};
const mockTab = { id: 123 };
(browser.tabs.create as jest.Mock).mockResolvedValue(mockTab);
// Act
await messageHandler.handleDecodedMessage(request);
// Assert
expect(browser.tabs.create).toHaveBeenCalledWith({
url: "https://allowed.com",
});
expect(mockClient.sendResourceToServer).toHaveBeenCalledWith({
resource: "opened-tab-id",
correlationId: "test-correlation-id",
tabId: 123,
});
});
});
describe("close-tabs command", () => {
it("should close tabs and send confirmation to the server", async () => {
// Arrange
const request: ServerMessageRequest = {
cmd: "close-tabs",
tabIds: [123, 456],
correlationId: "test-correlation-id",
};
(browser.tabs.remove as jest.Mock).mockResolvedValue(undefined);
// Act
await messageHandler.handleDecodedMessage(request);
// Assert
expect(browser.tabs.remove).toHaveBeenCalledWith([123, 456]);
expect(mockClient.sendResourceToServer).toHaveBeenCalledWith({
resource: "tabs-closed",
correlationId: "test-correlation-id",
});
});
});
describe("get-tab-list command", () => {
it("should get tabs and send them to the server", async () => {
// Arrange
const request: ServerMessageRequest = {
cmd: "get-tab-list",
correlationId: "test-correlation-id",
};
const mockTabs = [{ id: 123, url: "https://example.com" }];
(browser.tabs.query as jest.Mock).mockResolvedValue(mockTabs);
// Act
await messageHandler.handleDecodedMessage(request);
// Assert
expect(browser.tabs.query).toHaveBeenCalledWith({});
expect(mockClient.sendResourceToServer).toHaveBeenCalledWith({
resource: "tabs",
correlationId: "test-correlation-id",
tabs: mockTabs,
});
});
});
describe("get-browser-recent-history command", () => {
it("should get history items and send them to the server", async () => {
// Arrange
const request: ServerMessageRequest = {
cmd: "get-browser-recent-history",
searchQuery: "test",
correlationId: "test-correlation-id",
};
const mockHistoryItems = [
{ url: "https://example.com", title: "Example" },
{ url: "https://test.com", title: "Test" },
];
(browser.history.search as jest.Mock).mockResolvedValue(
mockHistoryItems
);
// Act
await messageHandler.handleDecodedMessage(request);
// Assert
expect(browser.history.search).toHaveBeenCalledWith({
text: "test",
maxResults: 200,
startTime: 0,
});
expect(mockClient.sendResourceToServer).toHaveBeenCalledWith({
resource: "history",
correlationId: "test-correlation-id",
historyItems: mockHistoryItems,
});
});
it("should use empty string for search query if not provided", async () => {
// Arrange
const request: ServerMessageRequest = {
cmd: "get-browser-recent-history",
correlationId: "test-correlation-id",
};
const mockHistoryItems = [
{ url: "https://example.com", title: "Example" },
];
(browser.history.search as jest.Mock).mockResolvedValue(
mockHistoryItems
);
// Act
await messageHandler.handleDecodedMessage(request);
// Assert
expect(browser.history.search).toHaveBeenCalledWith({
text: "",
maxResults: 200,
startTime: 0,
});
});
it("should filter out history items without URLs", async () => {
// Arrange
const request: ServerMessageRequest = {
cmd: "get-browser-recent-history",
correlationId: "test-correlation-id",
};
const mockHistoryItems = [
{ url: "https://example.com", title: "Example" },
{ title: "No URL" }, // This should be filtered out
];
(browser.history.search as jest.Mock).mockResolvedValue(
mockHistoryItems
);
// Act
await messageHandler.handleDecodedMessage(request);
// Assert
expect(mockClient.sendResourceToServer).toHaveBeenCalledWith({
resource: "history",
correlationId: "test-correlation-id",
historyItems: [{ url: "https://example.com", title: "Example" }],
});
});
});
describe("get-tab-content command", () => {
it("should get tab content and send it to the server", async () => {
// Arrange
const request: ServerMessageRequest = {
cmd: "get-tab-content",
tabId: 123,
correlationId: "test-correlation-id",
};
const mockTab = { id: 123, url: "https://example.com" };
(browser.tabs.get as jest.Mock).mockResolvedValue(mockTab);
(browser.permissions.contains as jest.Mock).mockResolvedValue(true);
const mockScriptResult = [
{
links: [{ url: "https://example.com/page", text: "Page" }],
fullText: "Page content",
isTruncated: false,
totalLength: 12,
},
];
(browser.tabs.executeScript as jest.Mock).mockResolvedValue(
mockScriptResult
);
// Act
await messageHandler.handleDecodedMessage(request);
// Assert
expect(browser.tabs.get).toHaveBeenCalledWith(123);
expect(browser.permissions.contains).toHaveBeenCalledWith({
origins: ["https://example.com/*"],
});
expect(browser.tabs.executeScript).toHaveBeenCalled();
expect(mockClient.sendResourceToServer).toHaveBeenCalledWith({
resource: "tab-content",
tabId: 123,
correlationId: "test-correlation-id",
isTruncated: false,
fullText: "Page content",
links: [{ url: "https://example.com/page", text: "Page" }],
totalLength: 12,
});
});
it("should throw an error if tab URL domain is in deny list", async () => {
// Arrange
const configWithDenyList: ExtensionConfig = {
secret: "test-secret",
toolSettings: {
"open-browser-tab": true,
"close-browser-tabs": true,
"get-list-of-open-tabs": true,
"get-recent-browser-history": true,
"get-tab-web-content": true,
"reorder-browser-tabs": true,
"find-highlight-in-browser-tab": true,
},
domainDenyList: ["example.com"], // Add example.com to deny list
ports: [8089],
auditLog: [],
};
(browser.storage.local.get as jest.Mock).mockResolvedValue({
config: configWithDenyList,
});
const request: ServerMessageRequest = {
cmd: "get-tab-content",
tabId: 123,
correlationId: "test-correlation-id",
};
const mockTab = { id: 123, url: "https://example.com" };
(browser.tabs.get as jest.Mock).mockResolvedValue(mockTab);
// Act & Assert
await expect(
messageHandler.handleDecodedMessage(request)
).rejects.toThrow("Domain in tab URL is in the deny list");
expect(browser.tabs.executeScript).not.toHaveBeenCalled();
});
it("should throw an error if permissions are denied", async () => {
// Arrange
const request: ServerMessageRequest = {
cmd: "get-tab-content",
tabId: 123,
correlationId: "test-correlation-id",
};
const mockTab = { id: 123, url: "https://example.com" };
(browser.tabs.get as jest.Mock).mockResolvedValue(mockTab);
(browser.permissions.contains as jest.Mock).mockResolvedValue(false);
// Act & Assert
await expect(
messageHandler.handleDecodedMessage(request)
).rejects.toThrow();
expect(browser.tabs.executeScript).not.toHaveBeenCalled();
});
});
describe("reorder-tabs command", () => {
it("should reorder tabs and send confirmation to the server", async () => {
// Arrange
const request: ServerMessageRequest = {
cmd: "reorder-tabs",
tabOrder: [123, 456, 789],
correlationId: "test-correlation-id",
};
(browser.tabs.move as jest.Mock).mockResolvedValue(undefined);
// Act
await messageHandler.handleDecodedMessage(request);
// Assert
expect(browser.tabs.move).toHaveBeenCalledTimes(3);
expect(browser.tabs.move).toHaveBeenNthCalledWith(1, 123, { index: 0 });
expect(browser.tabs.move).toHaveBeenNthCalledWith(2, 456, { index: 1 });
expect(browser.tabs.move).toHaveBeenNthCalledWith(3, 789, { index: 2 });
expect(mockClient.sendResourceToServer).toHaveBeenCalledWith({
resource: "tabs-reordered",
correlationId: "test-correlation-id",
tabOrder: [123, 456, 789],
});
});
});
describe("find-highlight command", () => {
it("should find and highlight text in a tab", async () => {
// Arrange
const request: ServerMessageRequest = {
cmd: "find-highlight",
tabId: 123,
queryPhrase: "test",
correlationId: "test-correlation-id",
};
const mockFindResults = { count: 5 };
(browser.find.find as jest.Mock).mockResolvedValue(mockFindResults);
(browser.tabs.update as jest.Mock).mockResolvedValue(undefined);
(browser.permissions.contains as jest.Mock).mockResolvedValue(true);
// Act
await messageHandler.handleDecodedMessage(request);
// Assert
expect(browser.find.find).toHaveBeenCalledWith("test", {
tabId: 123,
caseSensitive: true,
});
expect(browser.tabs.update).toHaveBeenCalledWith(123, { active: true });
expect(browser.find.highlightResults).toHaveBeenCalledWith({
tabId: 123,
});
expect(mockClient.sendResourceToServer).toHaveBeenCalledWith({
resource: "find-highlight-result",
correlationId: "test-correlation-id",
noOfResults: 5,
});
});
it("should not highlight or activate tab if no results found", async () => {
// Arrange
const request: ServerMessageRequest = {
cmd: "find-highlight",
tabId: 123,
queryPhrase: "test",
correlationId: "test-correlation-id",
};
const mockFindResults = { count: 0 };
const mockTab = { id: 123, url: "https://example.com" };
(browser.tabs.get as jest.Mock).mockResolvedValue(mockTab);
(browser.find.find as jest.Mock).mockResolvedValue(mockFindResults);
(browser.permissions.contains as jest.Mock).mockResolvedValue(true);
// Act
await messageHandler.handleDecodedMessage(request);
// Assert
expect(browser.tabs.update).not.toHaveBeenCalled();
expect(browser.find.highlightResults).not.toHaveBeenCalled();
expect(mockClient.sendResourceToServer).toHaveBeenCalledWith({
resource: "find-highlight-result",
correlationId: "test-correlation-id",
noOfResults: 0,
});
});
it("should throw an error if permissions are denied", async () => {
// Arrange
const request: ServerMessageRequest = {
cmd: "find-highlight",
tabId: 123,
queryPhrase: "test",
correlationId: "test-correlation-id",
};
const mockTab = { id: 123, url: "https://example.com" };
(browser.tabs.get as jest.Mock).mockResolvedValue(mockTab);
(browser.permissions.contains as jest.Mock).mockResolvedValue(false);
// Act & Assert
await expect(
messageHandler.handleDecodedMessage(request)
).rejects.toThrow();
expect(browser.find.find).not.toHaveBeenCalled();
});
});
});
});