|
1 | 1 | /** |
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 |
7 | 3 | */ |
8 | 4 |
|
9 | 5 | import { describe, it, expect, vi, beforeEach } from 'vitest'; |
| 6 | +import { |
| 7 | + evaluatePrivilegedScriptTool, |
| 8 | + handleEvaluatePrivilegedScript, |
| 9 | + isLikelyStatement, |
| 10 | +} from '../../src/tools/privileged-context.js'; |
10 | 11 |
|
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 | + }); |
16 | 35 |
|
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); |
21 | 38 | }); |
22 | 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(); |
| 63 | + const mockSetContext = vi.fn(); |
| 64 | + const mockSwitchToWindow = vi.fn(); |
| 65 | + |
23 | 66 | beforeEach(() => { |
24 | 67 | 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 | | - })); |
48 | 68 | }); |
49 | 69 |
|
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' }); |
52 | 73 |
|
53 | | - await handleSelectPrivilegedContext({ contextId: 'chrome-context-id' }); |
| 74 | + expect(result.isError).toBe(true); |
| 75 | + }); |
54 | 76 |
|
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' }); |
57 | 79 |
|
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); |
62 | 106 |
|
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' }); |
66 | 108 |
|
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 | + }); |
69 | 112 |
|
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' }); |
73 | 115 |
|
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 | + }); |
76 | 119 |
|
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' }); |
79 | 122 |
|
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 | + }); |
82 | 126 | }); |
83 | 127 | }); |
0 commit comments