refactor: split file and project zipping into seperate functions

Hweinstock · Hweinstock · commit b81a7c952912 · 2025-03-24T21:58:14.000-04:00
diff --git a/packages/core/src/codewhisperer/commands/startSecurityScan.ts b/packages/core/src/codewhisperer/commands/startSecurityScan.ts
@@ -157,7 +157,10 @@ export async function startSecurityScan(
             })
         }
         const zipType = scope === CodeAnalysisScope.PROJECT ? 'project' : 'file'
-        const zipMetadata = await zipUtil.generateZipCodeScan(editor?.document.uri, zipType)
+        const zipMetadata =
+            zipType === 'file'
+                ? await zipUtil.generateZipCodeScanForFile(editor?.document.uri)
+                : await zipUtil.generateZipCodeScanForProject()
         const projectPaths = getWorkspacePaths()
 
         const contextTruncationStartTime = performance.now()
diff --git a/packages/core/src/codewhisperer/util/zipUtil.ts b/packages/core/src/codewhisperer/util/zipUtil.ts
@@ -335,34 +335,53 @@ export class ZipUtil {
         return this._zipDir
     }
 
-    public async generateZipCodeScan(
+    public async generateZipCodeScanForFile(
         uri: vscode.Uri | undefined,
-        zipType: ZipType,
         options?: GenerateZipOptions
     ): Promise<ZipMetadata> {
+        try {
+            const zipDirPath = this.getZipDirPath()
+            const zipFilePath = await this.zipFile(uri, options?.includeGitDiff)
+
+            if (!options?.silent) {
+                getLogger().debug(`Picked source files: [${[...this._pickedSourceFiles].join(', ')}]`)
+            }
+            return {
+                rootDir: zipDirPath,
+                zipFilePath: zipFilePath,
+                srcPayloadSizeInBytes: this._totalSize,
+                scannedFiles: new Set([...this._pickedSourceFiles, ...this._pickedBuildFiles]),
+                zipFileSizeInBytes: (await fs.stat(zipFilePath)).size,
+                buildPayloadSizeInBytes: this._totalBuildSize,
+                lines: this._totalLines,
+                language: this._language,
+            }
+        } catch (error) {
+            getLogger().error('Zip error caused by: %O', error)
+            throw error
+        }
+    }
+
+    public async generateZipCodeScanForProject(options?: GenerateZipOptions): Promise<ZipMetadata> {
         try {
             const zipDirPath = this.getZipDirPath()
             // We assume there is at least one workspace open.
             const workspaceFolders = [...(vscode.workspace.workspaceFolders ?? [])] as CurrentWsFolders
 
-            const zipFilePath =
-                zipType === 'file'
-                    ? await this.zipFile(uri, options?.includeGitDiff)
-                    : await this.zipProject(workspaceFolders, defaultExcludePatterns, {
-                          includeGitDiff: options?.includeGitDiff,
-                          includeNonWorkspaceFiles: true,
-                      })
+            const zipFilePath = await this.zipProject(workspaceFolders, defaultExcludePatterns, {
+                includeGitDiff: options?.includeGitDiff,
+                includeNonWorkspaceFiles: true,
+            })
 
             if (!options?.silent) {
                 getLogger().debug(`Picked source files: [${[...this._pickedSourceFiles].join(', ')}]`)
             }
-            const zipFileSize = (await fs.stat(zipFilePath)).size
             return {
                 rootDir: zipDirPath,
                 zipFilePath: zipFilePath,
                 srcPayloadSizeInBytes: this._totalSize,
                 scannedFiles: new Set([...this._pickedSourceFiles, ...this._pickedBuildFiles]),
-                zipFileSizeInBytes: zipFileSize,
+                zipFileSizeInBytes: (await fs.stat(zipFilePath)).size,
                 buildPayloadSizeInBytes: this._totalBuildSize,
                 lines: this._totalLines,
                 language: this._language,
@@ -409,13 +428,13 @@ export class ZipUtil {
                     projectPath,
                 }
             )
-            const zipFileSize = (await fs.stat(zipFilePath)).size
+
             return {
                 rootDir: zipDirPath,
                 zipFilePath: zipFilePath,
                 srcPayloadSizeInBytes: this._totalSize,
                 scannedFiles: new Set(this._pickedSourceFiles),
-                zipFileSizeInBytes: zipFileSize,
+                zipFileSizeInBytes: (await fs.stat(zipFilePath)).size,
                 buildPayloadSizeInBytes: this._totalBuildSize,
                 lines: this._totalLines,
                 language: this._language,
diff --git a/packages/core/src/test/codewhisperer/zipUtil.test.ts b/packages/core/src/test/codewhisperer/zipUtil.test.ts
@@ -68,7 +68,7 @@ describe('zipUtil', function () {
         })
 
         it('Should generate zip for file scan and return expected metadata', async function () {
-            const zipMetadata = await zipUtil.generateZipCodeScan(vscode.Uri.file(appCodePath), 'file', {
+            const zipMetadata = await zipUtil.generateZipCodeScanForFile(vscode.Uri.file(appCodePath), {
                 includeGitDiff: true,
             })
             assert.strictEqual(zipMetadata.lines, 49)
@@ -84,13 +84,13 @@ describe('zipUtil', function () {
             sinon.stub(ZipUtil, 'aboveByteLimit').returns(true)
 
             await assert.rejects(
-                () => zipUtil.generateZipCodeScan(vscode.Uri.file(appCodePath), 'file', { includeGitDiff: true }),
+                () => zipUtil.generateZipCodeScanForFile(vscode.Uri.file(appCodePath), { includeGitDiff: true }),
                 new ToolkitError(`Payload size limit reached`, { code: 'FileSizeExceeded' })
             )
         })
 
         it('Should generate zip for project scan and return expected metadata', async function () {
-            const zipMetadata = await zipUtil.generateZipCodeScan(vscode.Uri.file(appCodePath), 'project')
+            const zipMetadata = await zipUtil.generateZipCodeScanForProject()
             assert.ok(zipMetadata.lines > 0)
             assert.ok(zipMetadata.rootDir.includes(codeScanTruncDirPrefix))
             assert.ok(zipMetadata.srcPayloadSizeInBytes > 0)
@@ -104,7 +104,7 @@ describe('zipUtil', function () {
             sinon.stub(ZipUtil, 'aboveByteLimit').returns(true)
 
             await assert.rejects(
-                () => zipUtil.generateZipCodeScan(vscode.Uri.file(appCodePath), 'project'),
+                () => zipUtil.generateZipCodeScanForProject(),
                 new ToolkitError('Payload size limit reached', { code: 'ProjectSizeExceeded' })
             )
         })
@@ -113,31 +113,29 @@ describe('zipUtil', function () {
             sinon.stub(ZipUtil, 'aboveByteLimit').returns(true)
 
             await assert.rejects(
-                () => zipUtil.generateZipCodeScan(vscode.Uri.file(appCodePath), 'project'),
+                () => zipUtil.generateZipCodeScanForProject(),
                 new ToolkitError('Payload size limit reached', { code: 'ProjectSizeExceeded' })
             )
         })
 
         it('Should read file content instead of from disk if file is dirty', async function () {
-            const zipMetadata = await zipUtil.generateZipCodeScan(vscode.Uri.file(appCodePath), 'project')
+            const zipMetadata = await zipUtil.generateZipCodeScanForProject()
 
             const document = await vscode.workspace.openTextDocument(appCodePath)
             await vscode.window.showTextDocument(document)
             void vscode.window.activeTextEditor?.edit((editBuilder) => {
                 editBuilder.insert(new vscode.Position(0, 0), '// a comment\n')
             })
 
-            const zipMetadata2 = await new ZipUtil(CodeWhispererConstants.codeScanTruncDirPrefix).generateZipCodeScan(
-                vscode.Uri.file(appCodePath),
-                'project'
-            )
+            const zipMetadata2 = await new ZipUtil(
+                CodeWhispererConstants.codeScanTruncDirPrefix
+            ).generateZipCodeScanForProject()
             assert.equal(zipMetadata2.lines, zipMetadata.lines + 1)
         })
 
         it('should handle path with repeated project name for file scan', async function () {
-            const zipMetadata = await zipUtil.generateZipCodeScan(
-                vscode.Uri.file(appCodePathWithRepeatedProjectName),
-                'file'
+            const zipMetadata = await zipUtil.generateZipCodeScanForFile(
+                vscode.Uri.file(appCodePathWithRepeatedProjectName)
             )
 
             const zipFileData = await fs.readFileBytes(zipMetadata.zipFilePath)
@@ -147,10 +145,7 @@ describe('zipUtil', function () {
         })
 
         it('should handle path with repeated project name for project scan', async function () {
-            const zipMetadata = await zipUtil.generateZipCodeScan(
-                vscode.Uri.file(appCodePathWithRepeatedProjectName),
-                'project'
-            )
+            const zipMetadata = await zipUtil.generateZipCodeScanForProject()
 
             const zipFileData = await fs.readFileBytes(zipMetadata.zipFilePath)
             const zip = await JSZip.loadAsync(zipFileData)
diff --git a/packages/core/src/testE2E/codewhisperer/securityScan.test.ts b/packages/core/src/testE2E/codewhisperer/securityScan.test.ts
@@ -91,11 +91,10 @@ describe('CodeWhisperer security scan', async function () {
     async function securityJobSetup(editor: vscode.TextEditor) {
         const codeScanStartTime = performance.now()
         const zipUtil = new ZipUtil(CodeWhispererConstants.codeScanTruncDirPrefix)
-        const uri = editor.document.uri
 
         const projectPaths = getWorkspacePaths()
         const scope = CodeWhispererConstants.CodeAnalysisScope.PROJECT
-        const zipMetadata = await zipUtil.generateZipCodeScan(uri, 'project')
+        const zipMetadata = await zipUtil.generateZipCodeScanForProject()
         const codeScanName = randomUUID()
 
         let artifactMap
diff --git a/packages/core/src/testInteg/perf/startSecurityScan.test.ts b/packages/core/src/testInteg/perf/startSecurityScan.test.ts
@@ -102,7 +102,7 @@ describe('startSecurityScanPerformanceTest', function () {
                     assert.ok(setup.commandSpy.neverCalledWith('workbench.action.problems.focus'))
                     assert.ok(setup.securityScanRenderSpy.calledOnce)
 
-                    assert.ok(setup.zipSpy.generateZipCodeScan.calledOnce)
+                    assert.ok(setup.zipSpy.generateZipCodeScanForProject.calledOnce)
                     assert.ok(setup.zipSpy.removeTmpFiles.calledOnce)
                     assert.ok(
                         getFsCallsUpperBound(setup.fsSpy) <= 5,
