From 4df3106b5ad8c4f09facaab0e523d0b8fbbbd803 Mon Sep 17 00:00:00 2001
From: Hweinstock <42325418+Hweinstock@users.noreply.github.com>
Date: Mon, 14 Oct 2024 16:11:11 -0400
Subject: [PATCH] test(amazonq): performance test for registerNewFiles (#5727)

## Problem
Continuation of https://github.com/aws/aws-toolkit-vscode/pull/5670

## Solution
same as previous

---

<!--- REMINDER: Ensure that your PR meets the guidelines in
CONTRIBUTING.md -->

License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.
---
 .../amazonqFeatureDev/session/sessionState.ts |  2 +-
 .../testInteg/perf/registerNewFiles.test.ts   | 87 +++++++++++++++++++
 2 files changed, 88 insertions(+), 1 deletion(-)
 create mode 100644 packages/core/src/testInteg/perf/registerNewFiles.test.ts

diff --git a/packages/core/src/amazonqFeatureDev/session/sessionState.ts b/packages/core/src/amazonqFeatureDev/session/sessionState.ts
index 9c802811eb..566bb312c5 100644
--- a/packages/core/src/amazonqFeatureDev/session/sessionState.ts
+++ b/packages/core/src/amazonqFeatureDev/session/sessionState.ts
@@ -57,7 +57,7 @@ export class ConversationNotStartedState implements Omit<SessionState, 'uploadId
     }
 }
 
-function registerNewFiles(
+export function registerNewFiles(
     fs: VirtualFileSystem,
     newFileContents: NewFileZipContents[],
     uploadId: string,
diff --git a/packages/core/src/testInteg/perf/registerNewFiles.test.ts b/packages/core/src/testInteg/perf/registerNewFiles.test.ts
new file mode 100644
index 0000000000..d658bf39de
--- /dev/null
+++ b/packages/core/src/testInteg/perf/registerNewFiles.test.ts
@@ -0,0 +1,87 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import assert from 'assert'
+import * as vscode from 'vscode'
+import { NewFileInfo, NewFileZipContents, registerNewFiles } from '../../amazonqFeatureDev'
+import { performanceTest } from '../../shared/performance/performance'
+import { getTestWorkspaceFolder } from '../integrationTestsUtilities'
+import { VirtualFileSystem } from '../../shared'
+
+interface SetupResult {
+    workspace: vscode.WorkspaceFolder
+    fileContents: NewFileZipContents[]
+}
+
+function getFileContents(numFiles: number, fileSize: number): NewFileZipContents[] {
+    return Array.from({ length: numFiles }, (_, i) => {
+        return {
+            zipFilePath: `test-path-${i}`,
+            fileContent: 'x'.repeat(fileSize),
+        }
+    })
+}
+
+function performanceTestWrapper(label: string, numFiles: number, fileSize: number) {
+    const conversationId = 'test-conversation'
+    return performanceTest(
+        {
+            testRuns: 10,
+            linux: {
+                userCpuUsage: 200,
+                systemCpuUsage: 35,
+                heapTotal: 8,
+            },
+            darwin: {
+                userCpuUsage: 200,
+                systemCpuUsage: 35,
+                heapTotal: 8,
+            },
+            win32: {
+                userCpuUsage: 200,
+                systemCpuUsage: 35,
+                heapTotal: 8,
+            },
+        },
+        label,
+        function () {
+            return {
+                setup: async () => {
+                    const testWorkspaceUri = vscode.Uri.file(getTestWorkspaceFolder())
+                    const fileContents = getFileContents(numFiles, fileSize)
+                    return {
+                        workspace: {
+                            uri: testWorkspaceUri,
+                            name: 'test-workspace',
+                            index: 0,
+                        },
+                        fileContents: fileContents,
+                    }
+                },
+                execute: async (setup: SetupResult) => {
+                    return registerNewFiles(
+                        new VirtualFileSystem(),
+                        setup.fileContents,
+                        'test-upload-id',
+                        [setup.workspace],
+                        conversationId
+                    )
+                },
+                verify: async (_setup: SetupResult, result: NewFileInfo[]) => {
+                    assert.strictEqual(result.length, numFiles)
+                },
+            }
+        }
+    )
+}
+
+describe('registerNewFiles', function () {
+    describe('performance tests', function () {
+        performanceTestWrapper('1x10MB', 1, 10000)
+        performanceTestWrapper('10x1000B', 10, 1000)
+        performanceTestWrapper('100x100B', 100, 100)
+        performanceTestWrapper('1000x10B', 1000, 10)
+        performanceTestWrapper('10000x1B', 10000, 1)
+    })
+})