Skip to content

Commit be27d83

Browse files
zixlin7zixlin7
committed
test(amazonq): add unit tests
1 parent 51f75d9 commit be27d83

File tree

4 files changed

+316
-7
lines changed

4 files changed

+316
-7
lines changed

‎packages/amazonq/src/app/inline/completion.ts‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ import {
2020
logInlineCompletionSessionResultsNotificationType,
2121
LogInlineCompletionSessionResultsParams,
2222
} from '@aws/language-server-runtimes/protocol'
23-
import { Commands } from 'aws-core-vscode/shared'
2423
import { SessionManager } from './sessionManager'
2524
import { RecommendationService } from './recommendationService'
2625

@@ -129,7 +128,7 @@ export class InlineCompletionManager implements Disposable {
129128
To show prev. and next. recommendation we need to re-register a new provider with the previous or next item
130129
*/
131130

132-
const prevCommand = Commands.declare('editor.action.inlineSuggest.showPrevious', () => async () => {
131+
const prevCommandHandler = async () => {
133132
SessionManager.instance.decrementActiveIndex()
134133
await commands.executeCommand('editor.action.inlineSuggest.hide')
135134
this.disposable.dispose()
@@ -138,10 +137,10 @@ export class InlineCompletionManager implements Disposable {
138137
new AmazonQInlineCompletionItemProvider(this.languageClient, false)
139138
)
140139
await commands.executeCommand('editor.action.inlineSuggest.trigger')
141-
})
142-
prevCommand.register()
140+
}
141+
commands.registerCommand('editor.action.inlineSuggest.showPrevious', prevCommandHandler)
143142

144-
const nextCommand = Commands.declare('editor.action.inlineSuggest.showNext', () => async () => {
143+
const nextCommandHandler = async () => {
145144
SessionManager.instance.incrementActiveIndex()
146145
await commands.executeCommand('editor.action.inlineSuggest.hide')
147146
this.disposable.dispose()
@@ -150,8 +149,8 @@ export class InlineCompletionManager implements Disposable {
150149
new AmazonQInlineCompletionItemProvider(this.languageClient, false)
151150
)
152151
await commands.executeCommand('editor.action.inlineSuggest.trigger')
153-
})
154-
nextCommand.register()
152+
}
153+
commands.registerCommand('editor.action.inlineSuggest.showNext', nextCommandHandler)
155154
}
156155
}
157156

‎packages/amazonq/src/app/inline/sessionManager.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,4 +92,9 @@ export class SessionManager {
9292
}
9393
return items
9494
}
95+
96+
public clear() {
97+
this.activeSession = undefined
98+
this.activeIndex = 0
99+
}
95100
}

‎packages/amazonq/test/unit/amazonq/apps/inline/completion.test.ts‎

Lines changed: 187 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,187 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
import sinon from 'sinon'
6+
import { commands, languages } from 'vscode'
7+
import assert from 'assert'
8+
import { LanguageClient } from 'vscode-languageclient'
9+
import { AmazonQInlineCompletionItemProvider, InlineCompletionManager } from '../../../../../src/app/inline/completion'
10+
11+
describe('InlineCompletionManager', () => {
12+
let manager: InlineCompletionManager
13+
let languageClient: LanguageClient
14+
let sendNotificationStub: sinon.SinonStub
15+
let registerProviderStub: sinon.SinonStub
16+
let registerCommandStub: sinon.SinonStub
17+
let executeCommandStub: sinon.SinonStub
18+
let disposableStub: sinon.SinonStub
19+
let sandbox: sinon.SinonSandbox
20+
21+
beforeEach(() => {
22+
sandbox = sinon.createSandbox()
23+
24+
registerProviderStub = sandbox.stub(languages, 'registerInlineCompletionItemProvider')
25+
registerCommandStub = sandbox.stub(commands, 'registerCommand')
26+
executeCommandStub = sandbox.stub(commands, 'executeCommand')
27+
sendNotificationStub = sandbox.stub()
28+
29+
// Create mock disposable
30+
const mockDisposable = {
31+
dispose: sandbox.stub(),
32+
}
33+
disposableStub = mockDisposable.dispose
34+
registerProviderStub.returns(mockDisposable)
35+
36+
languageClient = {
37+
sendNotification: sendNotificationStub,
38+
} as unknown as LanguageClient
39+
40+
manager = new InlineCompletionManager(languageClient)
41+
})
42+
43+
afterEach(() => {
44+
sandbox.restore()
45+
})
46+
47+
describe('constructor', () => {
48+
it('should initialize with language client and register provider', () => {
49+
assert(registerProviderStub.calledOnce)
50+
assert(
51+
registerProviderStub.calledWith(
52+
sinon.match.any,
53+
sinon.match.instanceOf(AmazonQInlineCompletionItemProvider)
54+
)
55+
)
56+
})
57+
})
58+
59+
describe('dispose', () => {
60+
it('should dispose of the disposable', () => {
61+
manager.dispose()
62+
assert(disposableStub.calledOnce)
63+
})
64+
})
65+
66+
describe('registerInlineCompletion', () => {
67+
beforeEach(() => {
68+
manager.registerInlineCompletion()
69+
})
70+
71+
it('should register accept and reject commands', () => {
72+
assert(registerCommandStub.calledWith('aws.sample-vscode-ext-amazonq.accept'))
73+
assert(registerCommandStub.calledWith('aws.sample-vscode-ext-amazonq.reject'))
74+
})
75+
76+
describe('onInlineAcceptance', () => {
77+
it('should send notification and re-register provider on acceptance', async () => {
78+
// Get the acceptance handler
79+
const acceptanceHandler = registerCommandStub
80+
.getCalls()
81+
?.find((call) => call.args[0] === 'aws.sample-vscode-ext-amazonq.accept')?.args[1]
82+
83+
const sessionId = 'test-session'
84+
const itemId = 'test-item'
85+
const requestStartTime = Date.now() - 1000
86+
const firstCompletionDisplayLatency = 500
87+
88+
await acceptanceHandler(sessionId, itemId, requestStartTime, firstCompletionDisplayLatency)
89+
90+
assert(sendNotificationStub.calledOnce)
91+
assert(
92+
sendNotificationStub.calledWith(
93+
sinon.match.any,
94+
sinon.match({
95+
sessionId,
96+
completionSessionResult: {
97+
[itemId]: {
98+
seen: true,
99+
accepted: true,
100+
discarded: false,
101+
},
102+
},
103+
})
104+
)
105+
)
106+
107+
assert(disposableStub.calledOnce)
108+
assert(registerProviderStub.calledTwice) // Once in constructor, once after acceptance
109+
})
110+
})
111+
112+
describe('onInlineRejection', () => {
113+
it('should hide suggestion and send notification on rejection', async () => {
114+
// Get the rejection handler
115+
const rejectionHandler = registerCommandStub
116+
.getCalls()
117+
.find((call) => call.args[0] === 'aws.sample-vscode-ext-amazonq.reject')?.args[1]
118+
119+
const sessionId = 'test-session'
120+
const itemId = 'test-item'
121+
122+
await rejectionHandler(sessionId, itemId)
123+
124+
assert(executeCommandStub.calledWith('editor.action.inlineSuggest.hide'))
125+
assert(sendNotificationStub.calledOnce)
126+
assert(
127+
sendNotificationStub.calledWith(
128+
sinon.match.any,
129+
sinon.match({
130+
sessionId,
131+
completionSessionResult: {
132+
[itemId]: {
133+
seen: true,
134+
accepted: false,
135+
discarded: false,
136+
},
137+
},
138+
})
139+
)
140+
)
141+
142+
assert(disposableStub.calledOnce)
143+
assert(registerProviderStub.calledTwice) // Once in constructor, once after rejection
144+
})
145+
})
146+
147+
describe('previous command', () => {
148+
it('should register and handle previous command correctly', async () => {
149+
const prevCommandCall = registerCommandStub
150+
.getCalls()
151+
.find((call) => call.args[0] === 'editor.action.inlineSuggest.showPrevious')
152+
153+
assert(prevCommandCall, 'Previous command should be registered')
154+
155+
if (prevCommandCall) {
156+
const handler = prevCommandCall.args[1]
157+
await handler()
158+
159+
assert(executeCommandStub.calledWith('editor.action.inlineSuggest.hide'))
160+
assert(disposableStub.calledOnce)
161+
assert(registerProviderStub.calledTwice)
162+
assert(executeCommandStub.calledWith('editor.action.inlineSuggest.trigger'))
163+
}
164+
})
165+
})
166+
167+
describe('next command', () => {
168+
it('should register and handle next command correctly', async () => {
169+
const nextCommandCall = registerCommandStub
170+
.getCalls()
171+
.find((call) => call.args[0] === 'editor.action.inlineSuggest.showNext')
172+
173+
assert(nextCommandCall, 'Next command should be registered')
174+
175+
if (nextCommandCall) {
176+
const handler = nextCommandCall.args[1]
177+
await handler()
178+
179+
assert(executeCommandStub.calledWith('editor.action.inlineSuggest.hide'))
180+
assert(disposableStub.calledOnce)
181+
assert(registerProviderStub.calledTwice)
182+
assert(executeCommandStub.calledWith('editor.action.inlineSuggest.trigger'))
183+
}
184+
})
185+
})
186+
})
187+
})

‎packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
import sinon from 'sinon'
7+
import { LanguageClient } from 'vscode-languageclient'
8+
import { Position, CancellationToken, InlineCompletionItem } from 'vscode'
9+
import assert from 'assert'
10+
import { RecommendationService } from '../../../../../src/app/inline/recommendationService'
11+
import { SessionManager } from '../../../../../src/app/inline/sessionManager'
12+
import { createMockDocument } from 'aws-core-vscode/test'
13+
14+
describe('RecommendationService', () => {
15+
let languageClient: LanguageClient
16+
let sendRequestStub: sinon.SinonStub
17+
let sandbox: sinon.SinonSandbox
18+
const mockDocument = createMockDocument()
19+
const mockPosition = { line: 0, character: 0 } as Position
20+
const mockContext = { triggerKind: 1, selectedCompletionInfo: undefined }
21+
const mockToken = { isCancellationRequested: false } as CancellationToken
22+
const mockInlineCompletionItemOne = {
23+
insertText: 'ItemOne',
24+
} as InlineCompletionItem
25+
26+
const mockInlineCompletionItemTwo = {
27+
insertText: 'ItemTwo',
28+
} as InlineCompletionItem
29+
const mockPartialResultToken = 'some-random-token'
30+
const service = RecommendationService.instance
31+
32+
beforeEach(() => {
33+
sandbox = sinon.createSandbox()
34+
35+
sendRequestStub = sandbox.stub()
36+
37+
languageClient = {
38+
sendRequest: sendRequestStub,
39+
} as unknown as LanguageClient
40+
})
41+
42+
afterEach(() => {
43+
sandbox.restore()
44+
SessionManager.instance.clear()
45+
})
46+
47+
describe('getAllRecommendations', () => {
48+
it('should handle single request with no partial result token', async () => {
49+
const mockFirstResult = {
50+
sessionId: 'test-session',
51+
items: [mockInlineCompletionItemOne],
52+
partialResultToken: undefined,
53+
}
54+
55+
sendRequestStub.resolves(mockFirstResult)
56+
57+
await service.getAllRecommendations(languageClient, mockDocument, mockPosition, mockContext, mockToken)
58+
59+
// Verify sendRequest was called with correct parameters
60+
assert(sendRequestStub.calledOnce)
61+
const requestArgs = sendRequestStub.firstCall.args[1]
62+
assert.deepStrictEqual(requestArgs, {
63+
textDocument: {
64+
uri: 'file:///test.py',
65+
},
66+
position: mockPosition,
67+
context: mockContext,
68+
})
69+
70+
// Verify session management
71+
const items = SessionManager.instance.getActiveRecommendation()
72+
assert.deepStrictEqual(items, [mockInlineCompletionItemOne])
73+
})
74+
75+
it('should handle multiple request with partial result token', async () => {
76+
const mockFirstResult = {
77+
sessionId: 'test-session',
78+
items: [mockInlineCompletionItemOne],
79+
partialResultToken: mockPartialResultToken,
80+
}
81+
82+
const mockSecondResult = {
83+
sessionId: 'test-session',
84+
items: [mockInlineCompletionItemTwo],
85+
partialResultToken: undefined,
86+
}
87+
88+
sendRequestStub.onFirstCall().resolves(mockFirstResult)
89+
sendRequestStub.onSecondCall().resolves(mockSecondResult)
90+
91+
await service.getAllRecommendations(languageClient, mockDocument, mockPosition, mockContext, mockToken)
92+
93+
// Verify sendRequest was called with correct parameters
94+
assert(sendRequestStub.calledTwice)
95+
const firstRequestArgs = sendRequestStub.firstCall.args[1]
96+
const expectedRequestArgs = {
97+
textDocument: {
98+
uri: 'file:///test.py',
99+
},
100+
position: mockPosition,
101+
context: mockContext,
102+
}
103+
const secondRequestArgs = sendRequestStub.secondCall.args[1]
104+
assert.deepStrictEqual(firstRequestArgs, expectedRequestArgs)
105+
assert.deepStrictEqual(secondRequestArgs, {
106+
...expectedRequestArgs,
107+
partialResultToken: mockPartialResultToken,
108+
})
109+
110+
// Verify session management
111+
const items = SessionManager.instance.getActiveRecommendation()
112+
assert.deepStrictEqual(items, [mockInlineCompletionItemOne, { insertText: '1' } as InlineCompletionItem])
113+
SessionManager.instance.incrementActiveIndex()
114+
const items2 = SessionManager.instance.getActiveRecommendation()
115+
assert.deepStrictEqual(items2, [mockInlineCompletionItemTwo, { insertText: '1' } as InlineCompletionItem])
116+
})
117+
})
118+
})

0 commit comments

Comments
 (0)