|
| 1 | +import * as assert from 'assert' |
| 2 | +import * as sinon from 'sinon' |
| 3 | +import { TextDocument } from 'vscode-languageserver-textdocument' |
| 4 | +import { InlineCompletionHandler } from './inlineCompletionHandler' |
| 5 | +import { SessionManager } from '../session/sessionManager' |
| 6 | +import { CodePercentageTracker } from '../tracker/codePercentageTracker' |
| 7 | +import { RecentEditTracker } from '../tracker/codeEditTracker' |
| 8 | +import { CursorTracker } from '../tracker/cursorTracker' |
| 9 | +import { StreakTracker } from '../tracker/streakTracker' |
| 10 | +import { TelemetryService } from '../../../shared/telemetry/telemetryService' |
| 11 | +import { UserWrittenCodeTracker } from '../../../shared/userWrittenCodeTracker' |
| 12 | +import { InlineCompletionTriggerKind, CancellationToken } from '@aws/language-server-runtimes/server-interface' |
| 13 | +import { EMPTY_RESULT } from '../contants/constants' |
| 14 | +import * as IdleWorkspaceManagerModule from '../../workspaceContext/IdleWorkspaceManager' |
| 15 | +import * as telemetryModule from '../telemetry/telemetry' |
| 16 | + |
| 17 | +describe('InlineCompletionHandler', () => { |
| 18 | + const testDocument = TextDocument.create('file:///test.cs', 'csharp', 1, 'test content') |
| 19 | + |
| 20 | + const completionParams = { |
| 21 | + textDocument: { uri: testDocument.uri }, |
| 22 | + position: { line: 0, character: 0 }, |
| 23 | + context: { triggerKind: InlineCompletionTriggerKind.Invoked }, |
| 24 | + } |
| 25 | + |
| 26 | + let handler: InlineCompletionHandler |
| 27 | + let completionSessionManager: SessionManager |
| 28 | + let amazonQServiceManager: any |
| 29 | + let codePercentageTracker: sinon.SinonStubbedInstance<CodePercentageTracker> |
| 30 | + let userWrittenCodeTracker: sinon.SinonStubbedInstance<UserWrittenCodeTracker> |
| 31 | + let recentEditTracker: sinon.SinonStubbedInstance<RecentEditTracker> |
| 32 | + let cursorTracker: sinon.SinonStubbedInstance<CursorTracker> |
| 33 | + let streakTracker: sinon.SinonStubbedInstance<StreakTracker> |
| 34 | + let telemetryService: sinon.SinonStubbedInstance<TelemetryService> |
| 35 | + let lsp: any |
| 36 | + let telemetry: any |
| 37 | + let credentialsProvider: any |
| 38 | + let workspace: any |
| 39 | + let logging: any |
| 40 | + let getTextDocument: sinon.SinonStub |
| 41 | + |
| 42 | + beforeEach(() => { |
| 43 | + SessionManager.reset() |
| 44 | + completionSessionManager = SessionManager.getInstance('COMPLETIONS') |
| 45 | + |
| 46 | + amazonQServiceManager = { |
| 47 | + getCodewhispererService: sinon.stub(), |
| 48 | + getConfiguration: sinon.stub().returns({ inlineSuggestions: {} }), |
| 49 | + } |
| 50 | + codePercentageTracker = sinon.createStubInstance(CodePercentageTracker) |
| 51 | + userWrittenCodeTracker = sinon.createStubInstance(UserWrittenCodeTracker) |
| 52 | + recentEditTracker = sinon.createStubInstance(RecentEditTracker) |
| 53 | + cursorTracker = sinon.createStubInstance(CursorTracker) |
| 54 | + streakTracker = sinon.createStubInstance(StreakTracker) |
| 55 | + telemetryService = sinon.createStubInstance(TelemetryService) |
| 56 | + |
| 57 | + workspace = { getWorkspaceFolder: sinon.stub() } |
| 58 | + logging = { log: sinon.stub(), debug: sinon.stub() } |
| 59 | + getTextDocument = sinon.stub().resolves(testDocument) |
| 60 | + lsp = { getClientInitializeParams: sinon.stub() } as any |
| 61 | + telemetry = { emitMetric: sinon.stub() } as any |
| 62 | + credentialsProvider = { getConnectionMetadata: sinon.stub() } as any |
| 63 | + |
| 64 | + // Stub IdleWorkspaceManager and telemetry functions |
| 65 | + sinon.stub(IdleWorkspaceManagerModule.IdleWorkspaceManager, 'recordActivityTimestamp') |
| 66 | + sinon.stub(telemetryModule, 'emitServiceInvocationTelemetry') |
| 67 | + sinon.stub(telemetryModule, 'emitServiceInvocationFailure') |
| 68 | + sinon.stub(telemetryModule, 'emitUserTriggerDecisionTelemetry') |
| 69 | + |
| 70 | + handler = new InlineCompletionHandler( |
| 71 | + logging, |
| 72 | + workspace, |
| 73 | + amazonQServiceManager, |
| 74 | + completionSessionManager, |
| 75 | + codePercentageTracker, |
| 76 | + userWrittenCodeTracker, |
| 77 | + recentEditTracker, |
| 78 | + cursorTracker, |
| 79 | + streakTracker, |
| 80 | + telemetry, |
| 81 | + telemetryService, |
| 82 | + credentialsProvider, |
| 83 | + () => false, |
| 84 | + () => 1000, |
| 85 | + lsp |
| 86 | + ) |
| 87 | + }) |
| 88 | + |
| 89 | + afterEach(() => { |
| 90 | + sinon.restore() |
| 91 | + }) |
| 92 | + |
| 93 | + it('should return empty result when concurrent request is in progress', async () => { |
| 94 | + // Make handler busy |
| 95 | + handler['isOnInlineCompletionHandlerInProgress'] = true |
| 96 | + |
| 97 | + const result = await handler.onInlineCompletion(completionParams, CancellationToken.None) |
| 98 | + |
| 99 | + assert.deepEqual(result, EMPTY_RESULT) |
| 100 | + sinon.assert.calledWith(logging.log, 'Skip concurrent inline completion') |
| 101 | + }) |
| 102 | + |
| 103 | + it('should return empty result when service manager not initialized', async () => { |
| 104 | + handler = new InlineCompletionHandler( |
| 105 | + logging, |
| 106 | + workspace, |
| 107 | + null as any, |
| 108 | + completionSessionManager, |
| 109 | + codePercentageTracker, |
| 110 | + userWrittenCodeTracker, |
| 111 | + recentEditTracker, |
| 112 | + cursorTracker, |
| 113 | + streakTracker, |
| 114 | + { emitMetric: sinon.stub() } as any, |
| 115 | + telemetryService, |
| 116 | + { getConnectionMetadata: sinon.stub() } as any, |
| 117 | + () => false, |
| 118 | + () => 1000, |
| 119 | + { getClientInitializeParams: sinon.stub() } as any |
| 120 | + ) |
| 121 | + |
| 122 | + const result = await handler.onInlineCompletion(completionParams, CancellationToken.None) |
| 123 | + |
| 124 | + assert.deepEqual(result, EMPTY_RESULT) |
| 125 | + sinon.assert.calledWith(logging.log, 'Amazon Q Service Manager not initialized yet') |
| 126 | + }) |
| 127 | + |
| 128 | + it('should return empty result when text document not found', async () => { |
| 129 | + getTextDocument.resolves(null) |
| 130 | + |
| 131 | + const result = await handler.onInlineCompletion(completionParams, CancellationToken.None) |
| 132 | + |
| 133 | + assert.deepEqual(result, EMPTY_RESULT) |
| 134 | + sinon.assert.calledWith(logging.log, `textDocument [${testDocument.uri}] not found`) |
| 135 | + }) |
| 136 | + |
| 137 | + it('should track cursor position when cursor tracker available', async () => { |
| 138 | + getTextDocument.resolves(null) // Will return early |
| 139 | + |
| 140 | + await handler.onInlineCompletion(completionParams, CancellationToken.None) |
| 141 | + |
| 142 | + sinon.assert.calledWith(cursorTracker.trackPosition, testDocument.uri, completionParams.position) |
| 143 | + }) |
| 144 | +}) |
0 commit comments