From f78cce03474f28f7b2938a5c85bf60cfa23aaf84 Mon Sep 17 00:00:00 2001 From: Laxman Reddy <141967714+laileni-aws@users.noreply.github.com> Date: Tue, 18 Feb 2025 16:59:38 -0800 Subject: [PATCH] Adding sourcefile to the zip before traversing the project --- .../unit/codewhisperer/util/zipUtil.test.ts | 9 ++++-- .../commands/startTestGeneration.ts | 2 +- .../core/src/codewhisperer/util/zipUtil.ts | 30 +++++++++++++---- .../src/shared/utilities/workspaceUtils.ts | 32 +++++++++++++++++-- 4 files changed, 61 insertions(+), 12 deletions(-) diff --git a/packages/amazonq/test/unit/codewhisperer/util/zipUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/zipUtil.test.ts index 230e950b045..d6887866640 100644 --- 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 index abc8c8ec74c..4309ce8c35c 100644 --- 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 index 18379e051ad..852af84e57b 100644 --- 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() - 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, 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 { + public async generateZipTestGen( + projectPath: string, + filePath: string, + initialExecution: boolean + ): Promise { 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 index f624a3ad2e1..d4e89f3e083 100644 --- 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> = [] + let totalSizeBytes = 0 const workspaceFoldersMapping = getWorkspaceFoldersByPrefixes(workspaceFolders) const workspaceToPrefix = new Map( @@ -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