chore(amazonq): move getTextDocument to textDocumentUtils and add test (#2382)

Author: atonaamz

server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.test.ts

@@ -12,7 +12,7 @@ import { TestFeatures } from '@aws/language-server-runtimes/testing'
 import * as assert from 'assert'
 import { AWSError } from 'aws-sdk'
 import sinon, { StubbedInstance } from 'ts-sinon'
-import { CodeWhispererServer, CodewhispererServerFactory, getLanguageIdFromUri } from './codeWhispererServer'
+import { CodeWhispererServer, CodewhispererServerFactory } from './codeWhispererServer'
 import {
     CodeWhispererServiceBase,
     CodeWhispererServiceToken,
@@ -2463,54 +2463,6 @@ describe('CodeWhisperer Server', () => {
         })
     })
 
-    describe('getLanguageIdFromUri', () => {
-        it('should return python for notebook cell URIs', () => {
-            const uri = 'vscode-notebook-cell:/some/path/notebook.ipynb#cell1'
-            assert.strictEqual(getLanguageIdFromUri(uri), 'python')
-        })
-
-        it('should return abap for files with ABAP extensions', () => {
-            const uris = ['file:///path/to/file.asprog']
-
-            uris.forEach(uri => {
-                assert.strictEqual(getLanguageIdFromUri(uri), 'abap')
-            })
-        })
-
-        it('should return empty string for non-ABAP files', () => {
-            const uris = ['file:///path/to/file.js', 'file:///path/to/file.ts', 'file:///path/to/file.py']
-
-            uris.forEach(uri => {
-                assert.strictEqual(getLanguageIdFromUri(uri), '')
-            })
-        })
-
-        it('should return empty string for invalid URIs', () => {
-            const invalidUris = ['', 'invalid-uri', 'file:///']
-
-            invalidUris.forEach(uri => {
-                assert.strictEqual(getLanguageIdFromUri(uri), '')
-            })
-        })
-
-        it('should log errors when provided with a logging object', () => {
-            const mockLogger = {
-                log: sinon.spy(),
-            }
-
-            const invalidUri = {} as string // Force type error
-            getLanguageIdFromUri(invalidUri, mockLogger)
-
-            sinon.assert.calledOnce(mockLogger.log)
-            sinon.assert.calledWith(mockLogger.log, sinon.match(/Error parsing URI to determine language:.*/))
-        })
-
-        it('should handle URIs without extensions', () => {
-            const uri = 'file:///path/to/file'
-            assert.strictEqual(getLanguageIdFromUri(uri), '')
-        })
-    })
-
     describe('Dynamic Service Manager Selection', () => {
         it('should use Token service manager when not using IAM auth', async () => {
             // Create isolated stubs for this test only

server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts

@@ -27,9 +27,6 @@ import { EditCompletionHandler } from './handler/editCompletionHandler'
 import { InlineCompletionHandler } from './handler/inlineCompletionHandler'
 import { SessionResultsHandler } from './handler/sessionResultsHandler'
 import { isUsingIAMAuth } from '../../shared/utils'
-import { TextDocument } from '@aws/language-server-runtimes/server-interface'
-import { ABAP_EXTENSIONS } from './contants/constants'
-import { URI } from 'vscode-uri'
 
 export const CodewhispererServerFactory =
     (serviceManager: (credentialsProvider?: any) => AmazonQBaseServiceManager): Server =>
@@ -179,8 +176,7 @@ export const CodewhispererServerFactory =
                 credentialsProvider,
                 () => editsEnabled,
                 () => timeSinceLastUserModification,
-                lsp,
-                getTextDocument
+                lsp
             )
 
             sessionResultsHandler = new SessionResultsHandler(
@@ -302,34 +298,6 @@ export const CodewhispererServerFactory =
         }
     }
 
-export const getLanguageIdFromUri = (uri: string, logging?: any): string => {
-    try {
-        if (uri.startsWith('vscode-notebook-cell:')) {
-            // use python for now as lsp does not support JL cell language detection
-            return 'python'
-        }
-        const extension = uri.split('.').pop()?.toLowerCase()
-        return ABAP_EXTENSIONS.has(extension || '') ? 'abap' : ''
-    } catch (err) {
-        logging?.log(`Error parsing URI to determine language: ${uri}: ${err}`)
-        return ''
-    }
-}
-
-export const getTextDocument = async (uri: string, workspace: any, logging: any): Promise<TextDocument | undefined> => {
-    let textDocument = await workspace.getTextDocument(uri)
-    if (!textDocument) {
-        try {
-            const content = await workspace.fs.readFile(URI.parse(uri).fsPath)
-            const languageId = getLanguageIdFromUri(uri)
-            textDocument = TextDocument.create(uri, languageId, 0, content)
-        } catch (err) {
-            logging.log(`Unable to load from ${uri}: ${err}`)
-        }
-    }
-    return textDocument
-}
-
 // Dynamic service manager factory that detects auth type at runtime
 export const CodeWhispererServer = CodewhispererServerFactory((credentialsProvider?: any) => {
     return isUsingIAMAuth(credentialsProvider) ? getOrThrowBaseIAMServiceManager() : getOrThrowBaseTokenServiceManager()

server/aws-lsp-codewhisperer/src/language-server/inline-completion/handler/inlineCompletionHandler.test.ts

@@ -13,6 +13,7 @@ import { InlineCompletionTriggerKind, CancellationToken } from '@aws/language-se
 import { EMPTY_RESULT } from '../contants/constants'
 import * as IdleWorkspaceManagerModule from '../../workspaceContext/IdleWorkspaceManager'
 import * as telemetryModule from '../telemetry/telemetry'
+import * as textDocumentUtils from '../utils/textDocumentUtils'
 
 describe('InlineCompletionHandler', () => {
     const testDocument = TextDocument.create('file:///test.cs', 'csharp', 1, 'test content')
@@ -37,7 +38,7 @@ describe('InlineCompletionHandler', () => {
     let credentialsProvider: any
     let workspace: any
     let logging: any
-    let getTextDocument: sinon.SinonStub
+    let getTextDocumentStub: sinon.SinonStub
 
     beforeEach(() => {
         SessionManager.reset()
@@ -56,16 +57,16 @@ describe('InlineCompletionHandler', () => {
 
         workspace = { getWorkspaceFolder: sinon.stub() }
         logging = { log: sinon.stub(), debug: sinon.stub() }
-        getTextDocument = sinon.stub().resolves(testDocument)
         lsp = { getClientInitializeParams: sinon.stub() } as any
         telemetry = { emitMetric: sinon.stub() } as any
         credentialsProvider = { getConnectionMetadata: sinon.stub() } as any
 
-        // Stub IdleWorkspaceManager and telemetry functions
+        // Stub IdleWorkspaceManager, telemetry functions, and textDocumentUtils
         sinon.stub(IdleWorkspaceManagerModule.IdleWorkspaceManager, 'recordActivityTimestamp')
         sinon.stub(telemetryModule, 'emitServiceInvocationTelemetry')
         sinon.stub(telemetryModule, 'emitServiceInvocationFailure')
         sinon.stub(telemetryModule, 'emitUserTriggerDecisionTelemetry')
+        getTextDocumentStub = sinon.stub(textDocumentUtils, 'getTextDocument')
 
         handler = new InlineCompletionHandler(
             logging,
@@ -82,8 +83,7 @@ describe('InlineCompletionHandler', () => {
             credentialsProvider,
             () => false,
             () => 1000,
-            lsp,
-            getTextDocument
+            lsp
         )
     })
 
@@ -117,8 +117,7 @@ describe('InlineCompletionHandler', () => {
             { getConnectionMetadata: sinon.stub() } as any,
             () => false,
             () => 1000,
-            { getClientInitializeParams: sinon.stub() } as any,
-            getTextDocument
+            { getClientInitializeParams: sinon.stub() } as any
         )
 
         const result = await handler.onInlineCompletion(completionParams, CancellationToken.None)
@@ -128,7 +127,7 @@ describe('InlineCompletionHandler', () => {
     })
 
     it('should return empty result when text document not found', async () => {
-        getTextDocument.resolves(null)
+        getTextDocumentStub.resolves(null)
 
         const result = await handler.onInlineCompletion(completionParams, CancellationToken.None)
 
@@ -137,7 +136,7 @@ describe('InlineCompletionHandler', () => {
     })
 
     it('should track cursor position when cursor tracker available', async () => {
-        getTextDocument.resolves(null) // Will return early
+        getTextDocumentStub.resolves(null) // Will return early
 
         await handler.onInlineCompletion(completionParams, CancellationToken.None)
 

server/aws-lsp-codewhisperer/src/language-server/inline-completion/handler/inlineCompletionHandler.ts

@@ -48,6 +48,7 @@ import {
 import { EMPTY_RESULT } from '../contants/constants'
 import { IdleWorkspaceManager } from '../../workspaceContext/IdleWorkspaceManager'
 import { mergeSuggestionsWithRightContext } from '../utils/mergeRightUtils'
+import { getTextDocument } from '../utils/textDocumentUtils'
 
 export class InlineCompletionHandler {
     private isOnInlineCompletionHandlerInProgress = false
@@ -67,12 +68,7 @@ export class InlineCompletionHandler {
         private readonly credentialsProvider: CredentialsProvider,
         private readonly getEditsEnabled: () => boolean,
         private readonly getTimeSinceLastUserModification: () => number,
-        private readonly lsp: Lsp,
-        private readonly getTextDocument: (
-            uri: string,
-            workspace: any,
-            logging: any
-        ) => Promise<TextDocument | undefined>
+        private readonly lsp: Lsp
     ) {}
 
     async onInlineCompletion(
@@ -110,7 +106,7 @@ export class InlineCompletionHandler {
             if (this.cursorTracker) {
                 this.cursorTracker.trackPosition(params.textDocument.uri, params.position)
             }
-            const textDocument = await this.getTextDocument(params.textDocument.uri, this.workspace, this.logging)
+            const textDocument = await getTextDocument(params.textDocument.uri, this.workspace, this.logging)
 
             const codeWhispererService = this.amazonQServiceManager.getCodewhispererService()
             const authType = codeWhispererService instanceof CodeWhispererServiceToken ? 'token' : 'iam'

server/aws-lsp-codewhisperer/src/language-server/inline-completion/utils/textDocumentUtils.test.ts

@@ -0,0 +1,113 @@
+import assert = require('assert')
+import { getLanguageIdFromUri, getTextDocument } from './textDocumentUtils'
+import { TextDocument } from '@aws/language-server-runtimes/server-interface'
+import sinon from 'ts-sinon'
+
+describe('textDocumentUtils', () => {
+    describe('getLanguageIdFromUri', () => {
+        it('should return python for notebook cell URIs', () => {
+            const uri = 'vscode-notebook-cell:/some/path/notebook.ipynb#cell1'
+            assert.strictEqual(getLanguageIdFromUri(uri), 'python')
+        })
+
+        it('should return abap for files with ABAP extensions', () => {
+            const uris = ['file:///path/to/file.asprog']
+
+            uris.forEach(uri => {
+                assert.strictEqual(getLanguageIdFromUri(uri), 'abap')
+            })
+        })
+
+        it('should return empty string for non-ABAP files', () => {
+            const uris = ['file:///path/to/file.js', 'file:///path/to/file.ts', 'file:///path/to/file.py']
+
+            uris.forEach(uri => {
+                assert.strictEqual(getLanguageIdFromUri(uri), '')
+            })
+        })
+
+        it('should return empty string for invalid URIs', () => {
+            const invalidUris = ['', 'invalid-uri', 'file:///']
+
+            invalidUris.forEach(uri => {
+                assert.strictEqual(getLanguageIdFromUri(uri), '')
+            })
+        })
+
+        it('should log errors when provided with a logging object', () => {
+            const mockLogger = {
+                log: sinon.spy(),
+            }
+
+            const invalidUri = {} as string // Force type error
+            getLanguageIdFromUri(invalidUri, mockLogger)
+
+            sinon.assert.calledOnce(mockLogger.log)
+            sinon.assert.calledWith(mockLogger.log, sinon.match(/Error parsing URI to determine language:.*/))
+        })
+
+        it('should handle URIs without extensions', () => {
+            const uri = 'file:///path/to/file'
+            assert.strictEqual(getLanguageIdFromUri(uri), '')
+        })
+    })
+
+    describe('getTextDocument', () => {
+        let mockWorkspace: any
+        let mockLogging: any
+
+        beforeEach(() => {
+            mockWorkspace = {
+                getTextDocument: sinon.stub(),
+                fs: {
+                    readFile: sinon.stub(),
+                },
+            }
+            mockLogging = {
+                log: sinon.stub(),
+            }
+        })
+
+        it('should return existing text document from workspace', async () => {
+            const existingDoc = TextDocument.create('file:///test.js', 'javascript', 1, 'content')
+            mockWorkspace.getTextDocument.resolves(existingDoc)
+
+            const result = await getTextDocument('file:///test.js', mockWorkspace, mockLogging)
+
+            assert.strictEqual(result, existingDoc)
+            sinon.assert.calledOnceWithExactly(mockWorkspace.getTextDocument, 'file:///test.js')
+            sinon.assert.notCalled(mockWorkspace.fs.readFile)
+        })
+
+        it('should create text document from file system when not in workspace', async () => {
+            mockWorkspace.getTextDocument.resolves(null)
+            mockWorkspace.fs.readFile.resolves('file content')
+
+            const result = await getTextDocument('file:///test.py', mockWorkspace, mockLogging)
+
+            assert.strictEqual(result?.uri, 'file:///test.py')
+            assert.strictEqual(result?.getText(), 'file content')
+            assert.strictEqual(result?.languageId, '')
+            sinon.assert.calledOnce(mockWorkspace.fs.readFile)
+        })
+
+        it('should handle file system read errors', async () => {
+            mockWorkspace.getTextDocument.resolves(null)
+            mockWorkspace.fs.readFile.rejects(new Error('File not found'))
+
+            const result = await getTextDocument('file:///missing.js', mockWorkspace, mockLogging)
+
+            assert.strictEqual(result, null)
+            sinon.assert.calledWith(mockLogging.log, sinon.match(/Unable to load from.*File not found/))
+        })
+
+        it('should use correct language ID for ABAP files', async () => {
+            mockWorkspace.getTextDocument.resolves(null)
+            mockWorkspace.fs.readFile.resolves('ABAP content')
+
+            const result = await getTextDocument('file:///test.asprog', mockWorkspace, mockLogging)
+
+            assert.strictEqual(result?.languageId, 'abap')
+        })
+    })
+})

server/aws-lsp-codewhisperer/src/language-server/inline-completion/utils/textDocumentUtils.ts

@@ -0,0 +1,31 @@
+import { TextDocument } from '@aws/language-server-runtimes/server-interface'
+import { ABAP_EXTENSIONS } from '../contants/constants'
+import { URI } from 'vscode-uri'
+
+export const getLanguageIdFromUri = (uri: string, logging?: any): string => {
+    try {
+        if (uri.startsWith('vscode-notebook-cell:')) {
+            // use python for now as lsp does not support JL cell language detection
+            return 'python'
+        }
+        const extension = uri.split('.').pop()?.toLowerCase()
+        return ABAP_EXTENSIONS.has(extension || '') ? 'abap' : ''
+    } catch (err) {
+        logging?.log(`Error parsing URI to determine language: ${uri}: ${err}`)
+        return ''
+    }
+}
+
+export const getTextDocument = async (uri: string, workspace: any, logging: any): Promise<TextDocument | undefined> => {
+    let textDocument = await workspace.getTextDocument(uri)
+    if (!textDocument) {
+        try {
+            const content = await workspace.fs.readFile(URI.parse(uri).fsPath)
+            const languageId = getLanguageIdFromUri(uri)
+            textDocument = TextDocument.create(uri, languageId, 0, content)
+        } catch (err) {
+            logging.log(`Unable to load from ${uri}: ${err}`)
+        }
+    }
+    return textDocument
+}