diff --git a/packages/core/src/testInteg/perf/downloadExportResultArchive.test.ts b/packages/core/src/testInteg/perf/downloadExportResultArchive.test.ts new file mode 100644 index 0000000000..4f78fdc4b0 --- /dev/null +++ b/packages/core/src/testInteg/perf/downloadExportResultArchive.test.ts @@ -0,0 +1,98 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +import assert from 'assert' +import { WorkspaceFolder } from 'vscode' +import { ExportResultArchiveCommandInput } from '@amzn/codewhisperer-streaming' +import * as sinon from 'sinon' +import path from 'path' +import { fs, getRandomString } from '../../shared' +import { createTestWorkspace } from '../../test/testUtil' +import { performanceTest } from '../../shared/performance/performance' +import { downloadExportResultArchive } from '../../shared/utilities/download' + +interface SetupResult { + workspace: WorkspaceFolder + exportCommandInput: ExportResultArchiveCommandInput + writeFileStub: sinon.SinonStub + cwStreaming: any +} + +interface FakeCommandOutput { + body: { binaryPayloadEvent: { bytes: Buffer } }[] +} + +function generateCommandOutput(numPieces: number, pieceSize: number): FakeCommandOutput { + const body = Array.from({ length: numPieces }, (_, i) => { + return { + binaryPayloadEvent: { + bytes: Buffer.from(getRandomString(pieceSize)), + }, + } + }) + return { body } +} + +async function setup(pieces: number, pieceSize: number): Promise { + // Force VSCode to find test workspace only to keep test contained and controlled. + const workspace = await createTestWorkspace(1, {}) + const exportCommandInput = {} as ExportResultArchiveCommandInput + // Manutally stub the CodeWhispererStreaming to avoid constructor call. + const cwStreaming = { exportResultArchive: () => generateCommandOutput(pieces, pieceSize) } + + const writeFileStub = sinon.stub(fs, 'writeFile') + return { workspace, exportCommandInput, writeFileStub, cwStreaming } +} + +function perfTest(pieces: number, pieceSize: number, label: string) { + return performanceTest( + { + testRuns: 10, + linux: { + userCpuUsage: 200, + systemCpuUsage: 35, + heapTotal: 4, + }, + darwin: { + userCpuUsage: 200, + systemCpuUsage: 35, + heapTotal: 4, + }, + win32: { + userCpuUsage: 200, + systemCpuUsage: 35, + heapTotal: 4, + }, + }, + label, + function () { + return { + setup: async () => await setup(pieces, pieceSize), + execute: async ({ workspace, cwStreaming, exportCommandInput, writeFileStub }: SetupResult) => { + await downloadExportResultArchive( + cwStreaming, + exportCommandInput, + path.join(workspace.uri.fsPath, 'result') + ) + }, + verify: async (setup: SetupResult) => { + assert.ok(setup.writeFileStub.calledOnce) + assert.ok((setup.writeFileStub.firstCall.args[1] as Buffer).length === pieces * pieceSize) + }, + } + } + ) +} + +describe('downloadExportResultArchive', function () { + describe('performanceTests', function () { + afterEach(function () { + sinon.restore() + }) + perfTest(1, 1000, '1x1KB') + perfTest(10, 100, '10x100B') + perfTest(100, 10, '100x10B') + perfTest(1000, 1, '1000x1B') + }) +})