Adding sourcefile to the zip before traversing the project

laileni-aws · Feb 19, 2025 · f78cce0 · f78cce0
1 parent 6d3670a
commit f78cce0
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 12 deletions.
diff --git a/packages/amazonq/test/unit/codewhisperer/util/zipUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/zipUtil.test.ts
@@ -159,7 +159,7 @@ describe('zipUtil', function () {
         it('should generate zip for test generation successfully', async function () {
             const mkdirSpy = sinon.spy(fs, 'mkdir')
 
-            const result = await zipUtil.generateZipTestGen(appRoot, false)
+            const result = await zipUtil.generateZipTestGen(appRoot, appCodePath, false)
 
             assert.ok(mkdirSpy.calledWith(path.join(testTempDirPath, 'utgRequiredArtifactsDir')))
             assert.ok(
@@ -184,7 +184,10 @@ describe('zipUtil', function () {
             }))
             sinon.stub(fs, 'mkdir').rejects(new Error('Directory creation failed'))
 
-            await assert.rejects(() => zipUtil.generateZipTestGen(appRoot, false), /Directory creation failed/)
+            await assert.rejects(
+                () => zipUtil.generateZipTestGen(appRoot, appCodePath, false),
+                /Directory creation failed/
+            )
         })
 
         it('Should handle zip project errors', async function () {
@@ -193,7 +196,7 @@ describe('zipUtil', function () {
             }))
             sinon.stub(zipUtil, 'zipProject' as keyof ZipUtil).rejects(new Error('Zip failed'))
 
-            await assert.rejects(() => zipUtil.generateZipTestGen(appRoot, false), /Zip failed/)
+            await assert.rejects(() => zipUtil.generateZipTestGen(appRoot, appCodePath, false), /Zip failed/)
         })
     })
 })
diff --git a/packages/core/src/codewhisperer/commands/startTestGeneration.ts b/packages/core/src/codewhisperer/commands/startTestGeneration.ts
@@ -58,7 +58,7 @@ export async function startTestGenerationProcess(
             session.projectRootPath = projectPath
             session.listOfTestGenerationJobId = []
         }
-        const zipMetadata = await zipUtil.generateZipTestGen(session.projectRootPath, initialExecution)
+        const zipMetadata = await zipUtil.generateZipTestGen(session.projectRootPath, filePath, initialExecution)
         session.srcPayloadSize = zipMetadata.buildPayloadSizeInBytes
         session.srcZipFileSize = zipMetadata.zipFileSizeInBytes
 

diff --git a/packages/core/src/codewhisperer/util/zipUtil.ts b/packages/core/src/codewhisperer/util/zipUtil.ts
@@ -204,7 +204,7 @@ export class ZipUtil {
         await processDirectory(metadataDir)
     }
 
-    protected async zipProject(useCase: FeatureUseCase, projectPath?: string, metadataDir?: string) {
+    protected async zipProject(useCase: FeatureUseCase, projectPath?: string, filePath?: string, metadataDir?: string) {
         const zip = new admZip()
         let projectPaths = []
         if (useCase === FeatureUseCase.TEST_GENERATION && projectPath) {
@@ -217,7 +217,13 @@ export class ZipUtil {
         }
         const languageCount = new Map<CodewhispererLanguage, number>()
 
-        await this.processSourceFiles(zip, languageCount, projectPaths, useCase)
+        await this.processSourceFiles(
+            zip,
+            languageCount,
+            projectPaths,
+            useCase,
+            useCase === FeatureUseCase.TEST_GENERATION ? filePath : undefined
+        )
         if (metadataDir) {
             await this.processMetadataDir(zip, metadataDir)
         }
@@ -404,7 +410,8 @@ export class ZipUtil {
         zip: admZip,
         languageCount: Map<CodewhispererLanguage, number>,
         projectPaths: string[] | undefined,
-        useCase: FeatureUseCase
+        useCase: FeatureUseCase,
+        filePath?: string
     ) {
         if (!projectPaths || projectPaths.length === 0) {
             return
@@ -414,7 +421,9 @@ export class ZipUtil {
             projectPaths,
             vscode.workspace.workspaceFolders as CurrentWsFolders,
             true,
-            this.getProjectScanPayloadSizeLimitInBytes()
+            this.getProjectScanPayloadSizeLimitInBytes(),
+            true,
+            useCase === FeatureUseCase.TEST_GENERATION ? filePath : undefined
         )
         for (const file of sourceFiles) {
             const projectName = path.basename(file.workspaceFolder.uri.fsPath)
@@ -560,7 +569,11 @@ export class ZipUtil {
         }
     }
 
-    public async generateZipTestGen(projectPath: string, initialExecution: boolean): Promise<ZipMetadata> {
+    public async generateZipTestGen(
+        projectPath: string,
+        filePath: string,
+        initialExecution: boolean
+    ): Promise<ZipMetadata> {
         try {
             // const repoMapFile = await LspClient.instance.getRepoMapJSON()
             const zipDirPath = this.getZipDirPath(FeatureUseCase.TEST_GENERATION)
@@ -591,7 +604,12 @@ export class ZipUtil {
                 }
             }
 
-            const zipFilePath: string = await this.zipProject(FeatureUseCase.TEST_GENERATION, projectPath, metadataDir)
+            const zipFilePath: string = await this.zipProject(
+                FeatureUseCase.TEST_GENERATION,
+                projectPath,
+                filePath,
+                metadataDir
+            )
             const zipFileSize = (await fs.stat(zipFilePath)).size
             return {
                 rootDir: zipDirPath,

diff --git a/packages/core/src/shared/utilities/workspaceUtils.ts b/packages/core/src/shared/utilities/workspaceUtils.ts
@@ -326,7 +326,8 @@ export async function collectFiles(
     workspaceFolders: CurrentWsFolders,
     respectGitIgnore: boolean = true,
     maxSize = 200 * 1024 * 1024, // 200 MB
-    defaultExcludePatterns: boolean = true
+    defaultExcludePatterns: boolean = true,
+    filePath?: string
 ): Promise<
     {
         workspaceFolder: vscode.WorkspaceFolder
@@ -337,6 +338,7 @@ export async function collectFiles(
     }[]
 > {
     const storage: Awaited<ReturnType<typeof collectFiles>> = []
+    let totalSizeBytes = 0
 
     const workspaceFoldersMapping = getWorkspaceFoldersByPrefixes(workspaceFolders)
     const workspaceToPrefix = new Map<vscode.WorkspaceFolder, string>(
@@ -359,7 +361,28 @@ export async function collectFiles(
         return prefix === '' ? path : `${prefix}/${path}`
     }
 
-    let totalSizeBytes = 0
+    // Add target file first to the collect files if provided
+    if (filePath) {
+        const targetUri = vscode.Uri.file(filePath)
+        const relativePath = getWorkspaceRelativePath(filePath, { workspaceFolders })
+
+        if (relativePath) {
+            const fileStat = await fs.stat(targetUri)
+
+            const fileContent = await readFile(targetUri)
+            if (fileContent !== undefined) {
+                totalSizeBytes += fileStat.size
+                storage.push({
+                    workspaceFolder: relativePath.workspaceFolder,
+                    relativeFilePath: relativePath.relativePath,
+                    fileUri: targetUri,
+                    fileContent: fileContent,
+                    zipFilePath: prefixWithFolderPrefix(relativePath.workspaceFolder, relativePath.relativePath),
+                })
+            }
+        }
+    }
+
     for (const rootPath of sourcePaths) {
         const allFiles = await vscode.workspace.findFiles(
             new vscode.RelativePattern(rootPath, '**'),
@@ -371,6 +394,11 @@ export async function collectFiles(
             : allFiles
 
         for (const file of files) {
+            // Skip if this file matches the target file that was already processed
+            if (filePath && file.fsPath === filePath) {
+                continue
+            }
+
             const relativePath = getWorkspaceRelativePath(file.fsPath, { workspaceFolders })
             if (!relativePath) {
                 continue