Skip to content

Commit 27de7db

Browse files
test: swap privileged-context.test.ts and privileged-context-state.test.ts
1 parent a5b5770 commit 27de7db

2 files changed

Lines changed: 164 additions & 164 deletions

File tree

Lines changed: 60 additions & 104 deletions
Original file line numberDiff line numberDiff line change
@@ -1,127 +1,83 @@
11
/**
2-
* Tests for statement detection and rejection in privileged context tools
2+
* Tests for privileged context state consistency
3+
*
4+
* Verifies that select_privileged_context updates currentContextId,
5+
* and that helper tools (set_firefox_prefs, list_extensions) don't
6+
* silently break a user's privileged context selection.
37
*/
48

59
import { describe, it, expect, vi, beforeEach } from 'vitest';
6-
import {
7-
evaluatePrivilegedScriptTool,
8-
handleEvaluatePrivilegedScript,
9-
isLikelyStatement,
10-
} from '../../src/tools/privileged-context.js';
1110

12-
// Mock the index module (used by handler tests)
13-
const mockGetFirefox = vi.hoisted(() => vi.fn());
14-
15-
vi.mock('../../src/index.js', () => ({
16-
getFirefox: () => mockGetFirefox(),
17-
}));
18-
19-
describe('Privileged Context Tool Definitions', () => {
20-
describe('evaluatePrivilegedScriptTool', () => {
21-
it('should have correct name', () => {
22-
expect(evaluatePrivilegedScriptTool.name).toBe('evaluate_privileged_script');
23-
});
24-
25-
it('should mention expression in description', () => {
26-
expect(evaluatePrivilegedScriptTool.description).toContain('expression');
27-
});
28-
});
29-
});
30-
31-
describe('isLikelyStatement', () => {
32-
it('should detect const declarations', () => {
33-
expect(isLikelyStatement('const x = 1')).toBe(true);
34-
});
35-
36-
it('should detect let declarations', () => {
37-
expect(isLikelyStatement('let x = 1')).toBe(true);
38-
});
39-
40-
it('should detect var declarations', () => {
41-
expect(isLikelyStatement('var x = 1')).toBe(true);
42-
});
43-
44-
it('should allow function calls', () => {
45-
expect(isLikelyStatement('Services.prefs.getBoolPref("foo")')).toBe(false);
46-
});
47-
48-
it('should allow simple expressions', () => {
49-
expect(isLikelyStatement('1 + 2')).toBe(false);
50-
});
51-
52-
it('should allow property access', () => {
53-
expect(isLikelyStatement('document.title')).toBe(false);
54-
});
55-
56-
it('should handle leading whitespace', () => {
57-
expect(isLikelyStatement(' const x = 1')).toBe(true);
58-
});
59-
});
60-
61-
describe('Privileged Context Tool Handlers', () => {
62-
const mockExecuteScript = vi.fn();
11+
describe('Privileged context state consistency', () => {
6312
const mockSetContext = vi.fn();
6413
const mockSwitchToWindow = vi.fn();
14+
const mockExecuteScript = vi.fn();
15+
const mockExecuteAsyncScript = vi.fn();
16+
17+
// Track currentContextId state as the real code would
18+
let mockCurrentContextId: string | null;
19+
const mockSetCurrentContextId = vi.fn((id: string) => {
20+
mockCurrentContextId = id;
21+
});
6522

6623
beforeEach(() => {
6724
vi.clearAllMocks();
25+
vi.resetModules();
26+
// Start with a content context (user has a normal tab open)
27+
mockCurrentContextId = 'original-content-context';
28+
29+
vi.doMock('../../src/index.js', () => ({
30+
getFirefox: vi.fn().mockResolvedValue({
31+
getDriver: () => ({
32+
switchTo: () => ({
33+
window: mockSwitchToWindow,
34+
}),
35+
setContext: mockSetContext,
36+
executeScript: mockExecuteScript,
37+
executeAsyncScript: mockExecuteAsyncScript,
38+
}),
39+
getCurrentContextId: () => mockCurrentContextId,
40+
setCurrentContextId: mockSetCurrentContextId,
41+
sendBiDiCommand: vi.fn().mockResolvedValue({
42+
contexts: [
43+
{ context: 'chrome-context-id', url: 'chrome://browser/content/browser.xhtml' },
44+
],
45+
}),
46+
}),
47+
}));
6848
});
6949

70-
describe('handleEvaluatePrivilegedScript', () => {
71-
it('should reject const statements with error', async () => {
72-
const result = await handleEvaluatePrivilegedScript({ expression: 'const x = 1' });
50+
it('select_privileged_context should update currentContextId', async () => {
51+
const { handleSelectPrivilegedContext } = await import('../../src/tools/privileged-context.js');
7352

74-
expect(result.isError).toBe(true);
75-
});
53+
await handleSelectPrivilegedContext({ contextId: 'chrome-context-id' });
7654

77-
it('should mention "statement" in error message', async () => {
78-
const result = await handleEvaluatePrivilegedScript({ expression: 'const x = 1' });
55+
expect(mockSwitchToWindow).toHaveBeenCalledWith('chrome-context-id');
56+
expect(mockSetContext).toHaveBeenCalledWith('chrome');
7957

80-
expect(result.content[0]).toHaveProperty('text', expect.stringMatching(/statement/i));
81-
});
82-
83-
it('should suggest IIFE workaround in error message', async () => {
84-
const result = await handleEvaluatePrivilegedScript({ expression: 'const x = 1' });
85-
86-
expect(result.content[0].text).toContain('function()');
87-
});
88-
89-
it('should return error when expression parameter is missing', async () => {
90-
const result = await handleEvaluatePrivilegedScript({});
91-
92-
expect(result.isError).toBe(true);
93-
expect(result.content[0].text).toContain('expression parameter is required');
94-
});
95-
96-
it('should execute valid expressions successfully', async () => {
97-
const mockFirefox = {
98-
getDriver: vi.fn().mockReturnValue({
99-
switchTo: () => ({ window: mockSwitchToWindow }),
100-
setContext: mockSetContext,
101-
executeScript: mockExecuteScript.mockResolvedValue('test-result'),
102-
}),
103-
};
104-
105-
mockGetFirefox.mockResolvedValue(mockFirefox);
58+
// select_privileged_context should call setCurrentContextId and update
59+
// currentContextId to 'chrome-context-id'
60+
expect(mockSetCurrentContextId).toHaveBeenCalledWith('chrome-context-id');
61+
});
10662

107-
const result = await handleEvaluatePrivilegedScript({ expression: 'document.title' });
63+
it('set_firefox_prefs after select_privileged_context should not revert to old context', async () => {
64+
const { handleSelectPrivilegedContext } = await import('../../src/tools/privileged-context.js');
65+
const { handleSetFirefoxPrefs } = await import('../../src/tools/firefox-prefs.js');
10866

109-
expect(result.isError).toBeUndefined();
110-
expect(result.content[0].text).toContain('test-result');
111-
});
67+
// User selects privileged context
68+
await handleSelectPrivilegedContext({ contextId: 'chrome-context-id' });
11269

113-
it('should reject let statements', async () => {
114-
const result = await handleEvaluatePrivilegedScript({ expression: 'let y = 2' });
70+
mockExecuteScript.mockResolvedValue(undefined);
71+
mockSwitchToWindow.mockClear();
72+
mockSetContext.mockClear();
11573

116-
expect(result.isError).toBe(true);
117-
expect(result.content[0]).toHaveProperty('text', expect.stringMatching(/statement/i));
118-
});
74+
// Call set_firefox_prefs which requires privileged context.
75+
await handleSetFirefoxPrefs({ prefs: { 'browser.ml.enable': true } });
11976

120-
it('should reject var statements', async () => {
121-
const result = await handleEvaluatePrivilegedScript({ expression: 'var z = 3' });
77+
const setContextCalls = mockSetContext.mock.calls;
78+
const lastSetContext = setContextCalls[setContextCalls.length - 1];
12279

123-
expect(result.isError).toBe(true);
124-
expect(result.content[0]).toHaveProperty('text', expect.stringMatching(/statement/i));
125-
});
80+
// Check that the context has not switched to content unexpectedly.
81+
expect(lastSetContext[0]).not.toBe('content');
12682
});
12783
});
Lines changed: 104 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -1,83 +1,127 @@
11
/**
2-
* Tests for privileged context state consistency
3-
*
4-
* Verifies that select_privileged_context updates currentContextId,
5-
* and that helper tools (set_firefox_prefs, list_extensions) don't
6-
* silently break a user's privileged context selection.
2+
* Tests for statement detection and rejection in privileged context tools
73
*/
84

95
import { describe, it, expect, vi, beforeEach } from 'vitest';
6+
import {
7+
evaluatePrivilegedScriptTool,
8+
handleEvaluatePrivilegedScript,
9+
isLikelyStatement,
10+
} from '../../src/tools/privileged-context.js';
1011

11-
describe('Privileged context state consistency', () => {
12-
const mockSetContext = vi.fn();
13-
const mockSwitchToWindow = vi.fn();
14-
const mockExecuteScript = vi.fn();
15-
const mockExecuteAsyncScript = vi.fn();
12+
// Mock the index module (used by handler tests)
13+
const mockGetFirefox = vi.hoisted(() => vi.fn());
14+
15+
vi.mock('../../src/index.js', () => ({
16+
getFirefox: () => mockGetFirefox(),
17+
}));
18+
19+
describe('Privileged Context Tool Definitions', () => {
20+
describe('evaluatePrivilegedScriptTool', () => {
21+
it('should have correct name', () => {
22+
expect(evaluatePrivilegedScriptTool.name).toBe('evaluate_privileged_script');
23+
});
24+
25+
it('should mention expression in description', () => {
26+
expect(evaluatePrivilegedScriptTool.description).toContain('expression');
27+
});
28+
});
29+
});
30+
31+
describe('isLikelyStatement', () => {
32+
it('should detect const declarations', () => {
33+
expect(isLikelyStatement('const x = 1')).toBe(true);
34+
});
1635

17-
// Track currentContextId state as the real code would
18-
let mockCurrentContextId: string | null;
19-
const mockSetCurrentContextId = vi.fn((id: string) => {
20-
mockCurrentContextId = id;
36+
it('should detect let declarations', () => {
37+
expect(isLikelyStatement('let x = 1')).toBe(true);
2138
});
2239

40+
it('should detect var declarations', () => {
41+
expect(isLikelyStatement('var x = 1')).toBe(true);
42+
});
43+
44+
it('should allow function calls', () => {
45+
expect(isLikelyStatement('Services.prefs.getBoolPref("foo")')).toBe(false);
46+
});
47+
48+
it('should allow simple expressions', () => {
49+
expect(isLikelyStatement('1 + 2')).toBe(false);
50+
});
51+
52+
it('should allow property access', () => {
53+
expect(isLikelyStatement('document.title')).toBe(false);
54+
});
55+
56+
it('should handle leading whitespace', () => {
57+
expect(isLikelyStatement(' const x = 1')).toBe(true);
58+
});
59+
});
60+
61+
describe('Privileged Context Tool Handlers', () => {
62+
const mockExecuteScript = vi.fn();
63+
const mockSetContext = vi.fn();
64+
const mockSwitchToWindow = vi.fn();
65+
2366
beforeEach(() => {
2467
vi.clearAllMocks();
25-
vi.resetModules();
26-
// Start with a content context (user has a normal tab open)
27-
mockCurrentContextId = 'original-content-context';
28-
29-
vi.doMock('../../src/index.js', () => ({
30-
getFirefox: vi.fn().mockResolvedValue({
31-
getDriver: () => ({
32-
switchTo: () => ({
33-
window: mockSwitchToWindow,
34-
}),
35-
setContext: mockSetContext,
36-
executeScript: mockExecuteScript,
37-
executeAsyncScript: mockExecuteAsyncScript,
38-
}),
39-
getCurrentContextId: () => mockCurrentContextId,
40-
setCurrentContextId: mockSetCurrentContextId,
41-
sendBiDiCommand: vi.fn().mockResolvedValue({
42-
contexts: [
43-
{ context: 'chrome-context-id', url: 'chrome://browser/content/browser.xhtml' },
44-
],
45-
}),
46-
}),
47-
}));
4868
});
4969

50-
it('select_privileged_context should update currentContextId', async () => {
51-
const { handleSelectPrivilegedContext } = await import('../../src/tools/privileged-context.js');
70+
describe('handleEvaluatePrivilegedScript', () => {
71+
it('should reject const statements with error', async () => {
72+
const result = await handleEvaluatePrivilegedScript({ expression: 'const x = 1' });
5273

53-
await handleSelectPrivilegedContext({ contextId: 'chrome-context-id' });
74+
expect(result.isError).toBe(true);
75+
});
5476

55-
expect(mockSwitchToWindow).toHaveBeenCalledWith('chrome-context-id');
56-
expect(mockSetContext).toHaveBeenCalledWith('chrome');
77+
it('should mention "statement" in error message', async () => {
78+
const result = await handleEvaluatePrivilegedScript({ expression: 'const x = 1' });
5779

58-
// select_privileged_context should call setCurrentContextId and update
59-
// currentContextId to 'chrome-context-id'
60-
expect(mockSetCurrentContextId).toHaveBeenCalledWith('chrome-context-id');
61-
});
80+
expect(result.content[0]).toHaveProperty('text', expect.stringMatching(/statement/i));
81+
});
82+
83+
it('should suggest IIFE workaround in error message', async () => {
84+
const result = await handleEvaluatePrivilegedScript({ expression: 'const x = 1' });
85+
86+
expect(result.content[0].text).toContain('function()');
87+
});
88+
89+
it('should return error when expression parameter is missing', async () => {
90+
const result = await handleEvaluatePrivilegedScript({});
91+
92+
expect(result.isError).toBe(true);
93+
expect(result.content[0].text).toContain('expression parameter is required');
94+
});
95+
96+
it('should execute valid expressions successfully', async () => {
97+
const mockFirefox = {
98+
getDriver: vi.fn().mockReturnValue({
99+
switchTo: () => ({ window: mockSwitchToWindow }),
100+
setContext: mockSetContext,
101+
executeScript: mockExecuteScript.mockResolvedValue('test-result'),
102+
}),
103+
};
104+
105+
mockGetFirefox.mockResolvedValue(mockFirefox);
62106

63-
it('set_firefox_prefs after select_privileged_context should not revert to old context', async () => {
64-
const { handleSelectPrivilegedContext } = await import('../../src/tools/privileged-context.js');
65-
const { handleSetFirefoxPrefs } = await import('../../src/tools/firefox-prefs.js');
107+
const result = await handleEvaluatePrivilegedScript({ expression: 'document.title' });
66108

67-
// User selects privileged context
68-
await handleSelectPrivilegedContext({ contextId: 'chrome-context-id' });
109+
expect(result.isError).toBeUndefined();
110+
expect(result.content[0].text).toContain('test-result');
111+
});
69112

70-
mockExecuteScript.mockResolvedValue(undefined);
71-
mockSwitchToWindow.mockClear();
72-
mockSetContext.mockClear();
113+
it('should reject let statements', async () => {
114+
const result = await handleEvaluatePrivilegedScript({ expression: 'let y = 2' });
73115

74-
// Call set_firefox_prefs which requires privileged context.
75-
await handleSetFirefoxPrefs({ prefs: { 'browser.ml.enable': true } });
116+
expect(result.isError).toBe(true);
117+
expect(result.content[0]).toHaveProperty('text', expect.stringMatching(/statement/i));
118+
});
76119

77-
const setContextCalls = mockSetContext.mock.calls;
78-
const lastSetContext = setContextCalls[setContextCalls.length - 1];
120+
it('should reject var statements', async () => {
121+
const result = await handleEvaluatePrivilegedScript({ expression: 'var z = 3' });
79122

80-
// Check that the context has not switched to content unexpectedly.
81-
expect(lastSetContext[0]).not.toBe('content');
123+
expect(result.isError).toBe(true);
124+
expect(result.content[0]).toHaveProperty('text', expect.stringMatching(/statement/i));
125+
});
82126
});
83127
});

0 commit comments

Comments
 (0)