diff --git a/packages/core/src/test/amazonqDoc/controller.test.ts b/packages/core/src/test/amazonqDoc/controller.test.ts index fda62e6fb55..1022900d331 100644 --- a/packages/core/src/test/amazonqDoc/controller.test.ts +++ b/packages/core/src/test/amazonqDoc/controller.test.ts @@ -24,305 +24,361 @@ import { FollowUpTypes } from '../../amazonq/commons/types' import { FileSystem } from '../../shared/fs/fs' import { ReadmeBuilder } from './mockContent' import * as path from 'path' - -describe('Controller - Doc Generation', () => { - const tabID = '123' - const conversationID = '456' - const uploadID = '789' - - let controllerSetup: ControllerSetup - let session: Session - let sendDocTelemetrySpy: sinon.SinonStub - let mockGetCodeGeneration: sinon.SinonStub - let getSessionStub: sinon.SinonStub - let modifiedReadme: string - const generatedReadme = ReadmeBuilder.createBaseReadme() - - const getFilePaths = (controllerSetup: ControllerSetup): NewFileInfo[] => [ - { - zipFilePath: path.normalize('README.md'), - relativePath: path.normalize('README.md'), - fileContent: generatedReadme, - rejected: false, - virtualMemoryUri: generateVirtualMemoryUri(uploadID, path.normalize('README.md'), docScheme), - workspaceFolder: controllerSetup.workspaceFolder, - changeApplied: false, - }, - ] - - async function createCodeGenState() { - mockGetCodeGeneration = sinon.stub().resolves({ codeGenerationStatus: { status: 'Complete' } }) - - const workspaceFolders = [controllerSetup.workspaceFolder] as CurrentWsFolders - const testConfig = { - conversationId: conversationID, - proxyClient: { - createConversation: () => sinon.stub(), - createUploadUrl: () => sinon.stub(), - generatePlan: () => sinon.stub(), - startCodeGeneration: () => sinon.stub(), - getCodeGeneration: () => mockGetCodeGeneration(), - exportResultArchive: () => sinon.stub(), - } as unknown as FeatureDevClient, - workspaceRoots: [''], - uploadId: uploadID, - workspaceFolders, - } - - const codeGenState = new DocCodeGenState(testConfig, getFilePaths(controllerSetup), [], [], tabID, 0, {}) - return createSession({ - messenger: controllerSetup.messenger, - sessionState: codeGenState, - conversationID, - tabID, - uploadID, - scheme: docScheme, - }) - } - async function fireFollowUps(followUpTypes: FollowUpTypes[]) { - for (const type of followUpTypes) { - controllerSetup.emitters.followUpClicked.fire({ +for (let i = 0; i < 1000; i++) { + describe(`Controller - Doc Generation ${i} times`, () => { + const tabID = '123' + const conversationID = '456' + const uploadID = '789' + + let controllerSetup: ControllerSetup + let session: Session + let sendDocTelemetrySpy: sinon.SinonStub + let mockGetCodeGeneration: sinon.SinonStub + let getSessionStub: sinon.SinonStub + let modifiedReadme: string + const generatedReadme = ReadmeBuilder.createBaseReadme() + let sandbox: sinon.SinonSandbox + + const getFilePaths = (controllerSetup: ControllerSetup): NewFileInfo[] => [ + { + zipFilePath: path.normalize('README.md'), + relativePath: path.normalize('README.md'), + fileContent: generatedReadme, + rejected: false, + virtualMemoryUri: generateVirtualMemoryUri(uploadID, path.normalize('README.md'), docScheme), + workspaceFolder: controllerSetup.workspaceFolder, + changeApplied: false, + }, + ] + + async function createCodeGenState(sandbox: sinon.SinonSandbox) { + mockGetCodeGeneration = sandbox.stub().resolves({ codeGenerationStatus: { status: 'Complete' } }) + + const workspaceFolders = [controllerSetup.workspaceFolder] as CurrentWsFolders + const testConfig = { + conversationId: conversationID, + proxyClient: { + createConversation: () => sandbox.stub(), + createUploadUrl: () => sandbox.stub(), + generatePlan: () => sandbox.stub(), + startCodeGeneration: () => sandbox.stub(), + getCodeGeneration: () => mockGetCodeGeneration(), + exportResultArchive: () => sandbox.stub(), + } as unknown as FeatureDevClient, + workspaceRoots: [''], + uploadId: uploadID, + workspaceFolders, + } + + const codeGenState = new DocCodeGenState(testConfig, getFilePaths(controllerSetup), [], [], tabID, 0, {}) + return createSession({ + messenger: controllerSetup.messenger, + sessionState: codeGenState, + conversationID, tabID, - followUp: { type }, + uploadID, + scheme: docScheme, + sandbox, }) } - } - - async function waitForStub(stub: sinon.SinonStub) { - await waitUntil(() => Promise.resolve(stub.callCount > 0), {}) - } - - async function performAction( - action: 'generate' | 'update' | 'makeChanges' | 'accept' | 'edit', - getSessionStub: sinon.SinonStub, - message?: string - ) { - const sequences = { - generate: FollowUpSequences.generateReadme, - update: FollowUpSequences.updateReadme, - edit: FollowUpSequences.editReadme, - makeChanges: FollowUpSequences.makeChanges, - accept: FollowUpSequences.acceptContent, + async function fireFollowUps(followUpTypes: FollowUpTypes[]) { + for (const type of followUpTypes) { + controllerSetup.emitters.followUpClicked.fire({ + tabID, + followUp: { type }, + }) + } } - await fireFollowUps(sequences[action]) - - if ((action === 'makeChanges' || action === 'edit') && message) { - controllerSetup.emitters.processHumanChatMessage.fire({ - tabID, - conversationID, - message, - }) + async function waitForStub(stub: sinon.SinonStub) { + await waitUntil(() => Promise.resolve(stub.callCount > 0), {}) } - await waitForStub(getSessionStub) - } - - before(() => { - sinon.stub(performance, 'now').returns(0) - }) - - beforeEach(async () => { - controllerSetup = await createController() - session = await createCodeGenState() - sendDocTelemetrySpy = sinon.stub(session, 'sendDocTelemetryEvent').resolves() - sinon.stub(session, 'preloader').resolves() - sinon.stub(session, 'send').resolves() - Object.defineProperty(session, '_conversationId', { - value: conversationID, - writable: true, - configurable: true, - }) - - sinon.stub(AuthUtil.instance, 'getChatAuthState').resolves({ - codewhispererCore: 'connected', - codewhispererChat: 'connected', - amazonQ: 'connected', - }) - sinon.stub(FileSystem.prototype, 'exists').resolves(false) - getSessionStub = sinon.stub(controllerSetup.sessionStorage, 'getSession').resolves(session) - modifiedReadme = ReadmeBuilder.createReadmeWithRepoStructure() - sinon - .stub(vscode.workspace, 'openTextDocument') - .callsFake(async (options?: string | vscode.Uri | { language?: string; content?: string }) => { - let documentPath = '' - if (typeof options === 'string') { - documentPath = options - } else if (options && 'path' in options) { - documentPath = options.path - } + async function performAction( + action: 'generate' | 'update' | 'makeChanges' | 'accept' | 'edit', + getSessionStub: sinon.SinonStub, + message?: string + ) { + const sequences = { + generate: FollowUpSequences.generateReadme, + update: FollowUpSequences.updateReadme, + edit: FollowUpSequences.editReadme, + makeChanges: FollowUpSequences.makeChanges, + accept: FollowUpSequences.acceptContent, + } + + await fireFollowUps(sequences[action]) + + if ((action === 'makeChanges' || action === 'edit') && message) { + controllerSetup.emitters.processHumanChatMessage.fire({ + tabID, + conversationID, + message, + }) + } + + await waitForStub(getSessionStub) + } - const isTempFile = documentPath === 'empty' - return { - getText: () => (isTempFile ? generatedReadme : modifiedReadme), - } as any + async function setupTest(sandbox: sinon.SinonSandbox) { + controllerSetup = await createController(sandbox) + session = await createCodeGenState(sandbox) + sendDocTelemetrySpy = sandbox.stub(session, 'sendDocTelemetryEvent').resolves() + sandbox.stub(session, 'preloader').resolves() + sandbox.stub(session, 'send').resolves() + Object.defineProperty(session, '_conversationId', { + value: conversationID, + writable: true, + configurable: true, }) - }) - afterEach(() => { - sinon.restore() - }) - - it('should emit generation telemetry for initial README generation', async () => { - await performAction('generate', getSessionStub) - - const expectedEvent = createExpectedEvent({ - type: 'generation', - ...EventMetrics.INITIAL_README, - interactionType: 'GENERATE_README', - conversationId: conversationID, - }) - - await assertTelemetry({ - spy: sendDocTelemetrySpy, - expectedEvent, - type: 'generation', - }) - }) - it('should emit another generation telemetry for make changes operation after initial README generation', async () => { - await performAction('generate', getSessionStub) - const firstExpectedEvent = createExpectedEvent({ - type: 'generation', - ...EventMetrics.INITIAL_README, - interactionType: 'GENERATE_README', - conversationId: conversationID, - }) - - await assertTelemetry({ - spy: sendDocTelemetrySpy, - expectedEvent: firstExpectedEvent, - type: 'generation', - }) - - await updateFilePaths(session, modifiedReadme, uploadID, docScheme, controllerSetup.workspaceFolder) - await performAction('makeChanges', getSessionStub, 'add repository structure section') - - const secondExpectedEvent = createExpectedEvent({ - type: 'generation', - ...EventMetrics.REPO_STRUCTURE, - interactionType: 'GENERATE_README', - conversationId: conversationID, - }) - await assertTelemetry({ - spy: sendDocTelemetrySpy, - expectedEvent: secondExpectedEvent, - type: 'generation', - callIndex: 1, - }) - }) + sandbox.stub(AuthUtil.instance, 'getChatAuthState').resolves({ + codewhispererCore: 'connected', + codewhispererChat: 'connected', + amazonQ: 'connected', + }) + sandbox.stub(FileSystem.prototype, 'exists').resolves(false) + getSessionStub = sandbox.stub(controllerSetup.sessionStorage, 'getSession').resolves(session) + modifiedReadme = ReadmeBuilder.createReadmeWithRepoStructure() + sandbox + .stub(vscode.workspace, 'openTextDocument') + .callsFake(async (options?: string | vscode.Uri | { language?: string; content?: string }) => { + let documentPath = '' + if (typeof options === 'string') { + documentPath = options + } else if (options && 'path' in options) { + documentPath = options.path + } + + const isTempFile = documentPath === 'empty' + return { + getText: () => (isTempFile ? generatedReadme : modifiedReadme), + } as any + }) + } - it('should emit acceptance telemetry for README generation', async () => { - await performAction('generate', getSessionStub) - await new Promise((resolve) => setTimeout(resolve, 100)) - const expectedEvent = createExpectedEvent({ - type: 'acceptance', - ...EventMetrics.INITIAL_README, - interactionType: 'GENERATE_README', - conversationId: conversationID, - }) + const retryTest = async ( + testMethod: () => Promise, + maxRetries: number = 3, + delayMs: number = 1000 + ): Promise => { + let lastError: Error | undefined + + for (let attempt = 1; attempt <= maxRetries + 1; attempt++) { + sandbox = sinon.createSandbox() + try { + await setupTest(sandbox) + await testMethod() + sandbox.restore() + return + } catch (error) { + lastError = error as Error + sandbox.restore() + + if (attempt > maxRetries) { + console.error(`Test failed after ${maxRetries} retries:`, lastError) + throw lastError + } + + console.log(`Test attempt ${attempt} failed, retrying...`) + await new Promise((resolve) => setTimeout(resolve, delayMs)) + } + } + } - await performAction('accept', getSessionStub) - await assertTelemetry({ - spy: sendDocTelemetrySpy, - expectedEvent, - type: 'acceptance', - callIndex: 1, - }) - }) - it('should emit generation telemetry for README update', async () => { - await performAction('update', getSessionStub) - - const expectedEvent = createExpectedEvent({ - type: 'generation', - ...EventMetrics.REPO_STRUCTURE, - interactionType: 'UPDATE_README', - conversationId: conversationID, + after(() => { + if (sandbox) { + sandbox.restore() + } }) - await assertTelemetry({ - spy: sendDocTelemetrySpy, - expectedEvent, - type: 'generation', + it('should emit generation telemetry for initial README generation', async () => { + await retryTest(async () => { + await performAction('generate', getSessionStub) + + const expectedEvent = createExpectedEvent({ + type: 'generation', + ...EventMetrics.INITIAL_README, + interactionType: 'GENERATE_README', + conversationId: conversationID, + }) + + await assertTelemetry({ + spy: sendDocTelemetrySpy, + expectedEvent, + type: 'generation', + sandbox, + }) + }) }) - }) - it('should emit another generation telemetry for make changes operation after README update', async () => { - await performAction('update', getSessionStub) - await new Promise((resolve) => setTimeout(resolve, 100)) - - modifiedReadme = ReadmeBuilder.createReadmeWithDataFlow() - await updateFilePaths(session, modifiedReadme, uploadID, docScheme, controllerSetup.workspaceFolder) - - await performAction('makeChanges', getSessionStub, 'add data flow section') - - const expectedEvent = createExpectedEvent({ - type: 'generation', - ...EventMetrics.DATA_FLOW, - interactionType: 'UPDATE_README', - conversationId: conversationID, - callIndex: 1, + it('should emit another generation telemetry for make changes operation after initial README generation', async () => { + await retryTest(async () => { + await performAction('generate', getSessionStub) + const firstExpectedEvent = createExpectedEvent({ + type: 'generation', + ...EventMetrics.INITIAL_README, + interactionType: 'GENERATE_README', + conversationId: conversationID, + }) + + await assertTelemetry({ + spy: sendDocTelemetrySpy, + expectedEvent: firstExpectedEvent, + type: 'generation', + sandbox, + }) + + await updateFilePaths(session, modifiedReadme, uploadID, docScheme, controllerSetup.workspaceFolder) + await performAction('makeChanges', getSessionStub, 'add repository structure section') + + const secondExpectedEvent = createExpectedEvent({ + type: 'generation', + ...EventMetrics.REPO_STRUCTURE, + interactionType: 'GENERATE_README', + conversationId: conversationID, + }) + + await assertTelemetry({ + spy: sendDocTelemetrySpy, + expectedEvent: secondExpectedEvent, + type: 'generation', + callIndex: 1, + sandbox, + }) + }) }) - await assertTelemetry({ - spy: sendDocTelemetrySpy, - expectedEvent, - type: 'generation', - callIndex: 1, + it('should emit acceptance telemetry for README generation', async () => { + await retryTest(async () => { + await performAction('generate', getSessionStub) + await new Promise((resolve) => setTimeout(resolve, 100)) + const expectedEvent = createExpectedEvent({ + type: 'acceptance', + ...EventMetrics.INITIAL_README, + interactionType: 'GENERATE_README', + conversationId: conversationID, + }) + + await performAction('accept', getSessionStub) + await assertTelemetry({ + spy: sendDocTelemetrySpy, + expectedEvent, + type: 'acceptance', + callIndex: 1, + sandbox, + }) + }) }) - }) - - it('should emit acceptance telemetry for README update', async () => { - await performAction('update', getSessionStub) - await new Promise((resolve) => setTimeout(resolve, 100)) - - const expectedEvent = createExpectedEvent({ - type: 'acceptance', - ...EventMetrics.REPO_STRUCTURE, - interactionType: 'UPDATE_README', - conversationId: conversationID, + it('should emit generation telemetry for README update', async () => { + await retryTest(async () => { + await performAction('update', getSessionStub) + + const expectedEvent = createExpectedEvent({ + type: 'generation', + ...EventMetrics.REPO_STRUCTURE, + interactionType: 'UPDATE_README', + conversationId: conversationID, + }) + + await assertTelemetry({ + spy: sendDocTelemetrySpy, + expectedEvent, + type: 'generation', + sandbox, + }) + }) }) - - await performAction('accept', getSessionStub) - await assertTelemetry({ - spy: sendDocTelemetrySpy, - expectedEvent, - type: 'acceptance', - callIndex: 1, + it('should emit another generation telemetry for make changes operation after README update', async () => { + await retryTest(async () => { + await performAction('update', getSessionStub) + await new Promise((resolve) => setTimeout(resolve, 100)) + + modifiedReadme = ReadmeBuilder.createReadmeWithDataFlow() + await updateFilePaths(session, modifiedReadme, uploadID, docScheme, controllerSetup.workspaceFolder) + + await performAction('makeChanges', getSessionStub, 'add data flow section') + + const expectedEvent = createExpectedEvent({ + type: 'generation', + ...EventMetrics.DATA_FLOW, + interactionType: 'UPDATE_README', + conversationId: conversationID, + callIndex: 1, + }) + + await assertTelemetry({ + spy: sendDocTelemetrySpy, + expectedEvent, + type: 'generation', + callIndex: 1, + sandbox, + }) + }) }) - }) - it('should emit generation telemetry for README edit', async () => { - await performAction('edit', getSessionStub, 'add repository structure section') - - const expectedEvent = createExpectedEvent({ - type: 'generation', - ...EventMetrics.REPO_STRUCTURE, - interactionType: 'EDIT_README', - conversationId: conversationID, + it('should emit acceptance telemetry for README update', async () => { + await retryTest(async () => { + await performAction('update', getSessionStub) + await new Promise((resolve) => setTimeout(resolve, 100)) + + const expectedEvent = createExpectedEvent({ + type: 'acceptance', + ...EventMetrics.REPO_STRUCTURE, + interactionType: 'UPDATE_README', + conversationId: conversationID, + }) + + await performAction('accept', getSessionStub) + await assertTelemetry({ + spy: sendDocTelemetrySpy, + expectedEvent, + type: 'acceptance', + callIndex: 1, + sandbox, + }) + }) }) - await assertTelemetry({ - spy: sendDocTelemetrySpy, - expectedEvent, - type: 'generation', - }) - }) - it('should emit acceptance telemetry for README edit', async () => { - await performAction('edit', getSessionStub, 'add repository structure section') - await new Promise((resolve) => setTimeout(resolve, 100)) - - const expectedEvent = createExpectedEvent({ - type: 'acceptance', - ...EventMetrics.REPO_STRUCTURE, - interactionType: 'EDIT_README', - conversationId: conversationID, + it('should emit generation telemetry for README edit', async () => { + await retryTest(async () => { + await performAction('edit', getSessionStub, 'add repository structure section') + + const expectedEvent = createExpectedEvent({ + type: 'generation', + ...EventMetrics.REPO_STRUCTURE, + interactionType: 'EDIT_README', + conversationId: conversationID, + }) + + await assertTelemetry({ + spy: sendDocTelemetrySpy, + expectedEvent, + type: 'generation', + sandbox, + }) + }) }) - - await performAction('accept', getSessionStub) - await assertTelemetry({ - spy: sendDocTelemetrySpy, - expectedEvent, - type: 'acceptance', - callIndex: 1, + it('should emit acceptance telemetry for README edit', async () => { + await retryTest(async () => { + await performAction('edit', getSessionStub, 'add repository structure section') + await new Promise((resolve) => setTimeout(resolve, 100)) + + const expectedEvent = createExpectedEvent({ + type: 'acceptance', + ...EventMetrics.REPO_STRUCTURE, + interactionType: 'EDIT_README', + conversationId: conversationID, + }) + + await performAction('accept', getSessionStub) + await assertTelemetry({ + spy: sendDocTelemetrySpy, + expectedEvent, + type: 'acceptance', + callIndex: 1, + sandbox, + }) + }) }) }) -}) +} diff --git a/packages/core/src/test/amazonqDoc/utils.ts b/packages/core/src/test/amazonqDoc/utils.ts index 8e0a54d8863..6e4c173f8fc 100644 --- a/packages/core/src/test/amazonqDoc/utils.ts +++ b/packages/core/src/test/amazonqDoc/utils.ts @@ -22,9 +22,9 @@ import { DocGenerationTask } from '../../amazonqDoc/controllers/docGenerationTas import { DocV2GenerationEvent, DocV2AcceptanceEvent } from '../../amazonqFeatureDev/client/featuredevproxyclient' import { FollowUpTypes } from '../../amazonq/commons/types' -export function createMessenger(): DocMessenger { +export function createMessenger(sandbox: sinon.SinonSandbox): DocMessenger { return new DocMessenger( - new AppToWebViewMessageDispatcher(new MessagePublisher(sinon.createStubInstance(vscode.EventEmitter))), + new AppToWebViewMessageDispatcher(new MessagePublisher(sandbox.createStubInstance(vscode.EventEmitter))), docChat ) } @@ -61,6 +61,7 @@ export async function createSession({ conversationID = '0', tabID = '0', uploadID = '0', + sandbox, }: { messenger: DocMessenger scheme: string @@ -68,19 +69,19 @@ export async function createSession({ conversationID?: string tabID?: string uploadID?: string + sandbox: sinon.SinonSandbox }) { const sessionConfig = await createSessionConfig(scheme) - const client = sinon.createStubInstance(FeatureDevClient) + const client = sandbox.createStubInstance(FeatureDevClient) client.createConversation.resolves(conversationID) const session = new Session(sessionConfig, messenger, tabID, sessionState, client) - sinon.stub(session, 'conversationId').get(() => conversationID) - sinon.stub(session, 'uploadId').get(() => uploadID) + sandbox.stub(session, 'conversationId').get(() => conversationID) + sandbox.stub(session, 'uploadId').get(() => uploadID) return session } - export async function sessionRegisterProvider(session: Session, uri: vscode.Uri, fileContents: Uint8Array) { session.config.fs.registerProvider(uri, new VirtualMemoryFile(fileContents)) } @@ -98,12 +99,12 @@ export async function sessionWriteFile(session: Session, uri: vscode.Uri, encode }) } -export async function createController(): Promise { - const messenger = createMessenger() +export async function createController(sandbox: sinon.SinonSandbox): Promise { + const messenger = createMessenger(sandbox) // Create a new workspace root const testWorkspaceFolder = await createTestWorkspaceFolder() - sinon.stub(vscode.workspace, 'workspaceFolders').value([testWorkspaceFolder]) + sandbox.stub(vscode.workspace, 'workspaceFolders').value([testWorkspaceFolder]) const sessionStorage = new DocChatSessionStorage(messenger) @@ -113,7 +114,7 @@ export async function createController(): Promise { mockChatControllerEventEmitters, messenger, sessionStorage, - sinon.createStubInstance(vscode.EventEmitter).event + sandbox.createStubInstance(vscode.EventEmitter).event ) new DocGenerationTask() @@ -199,10 +200,11 @@ export async function assertTelemetry(params: { expectedEvent: DocV2GenerationEvent | DocV2AcceptanceEvent type: 'generation' | 'acceptance' callIndex?: number + sandbox: sinon.SinonSandbox }) { await new Promise((resolve) => setTimeout(resolve, 100)) const spyCall = params.callIndex !== undefined ? params.spy.getCall(params.callIndex) : params.spy - sinon.assert.calledWith(spyCall, sinon.match(params.expectedEvent), params.type) + params.sandbox.assert.calledWith(spyCall, params.sandbox.match(params.expectedEvent), params.type) } export async function updateFilePaths(