test: improve doc telemetry test robustness

Author: KevinDing1

packages/core/src/test/amazonqDoc/controller.test.ts

@@ -14,6 +14,7 @@ import {
     FollowUpSequences,
     generateVirtualMemoryUri,
     updateFilePaths,
+    waitForTelemetryCall,
 } from './utils'
 import { CurrentWsFolders, NewFileInfo } from '../../amazonqDoc/types'
 import { CodeGenState, docScheme, Session } from '../../amazonqDoc'
@@ -165,6 +166,7 @@ for (let i = 0; i < 1000; i++) {
         it('should emit generation telemetry for initial README generation', async () => {
             await performAction('generate', getSessionStub)
 
+            await waitForTelemetryCall(sendDocTelemetrySpy, 1)
             const expectedEvent = createExpectedEvent({
                 type: 'generation',
                 ...EventMetrics.INITIAL_README,
@@ -186,7 +188,7 @@ for (let i = 0; i < 1000; i++) {
                 interactionType: 'GENERATE_README',
                 conversationId: conversationID,
             })
-
+            await waitForTelemetryCall(sendDocTelemetrySpy, 1)
             await assertTelemetry({
                 spy: sendDocTelemetrySpy,
                 expectedEvent: firstExpectedEvent,
@@ -202,7 +204,7 @@ for (let i = 0; i < 1000; i++) {
                 interactionType: 'GENERATE_README',
                 conversationId: conversationID,
             })
-
+            await waitForTelemetryCall(sendDocTelemetrySpy, 1)
             await assertTelemetry({
                 spy: sendDocTelemetrySpy,
                 expectedEvent: secondExpectedEvent,
@@ -213,15 +215,15 @@ for (let i = 0; i < 1000; i++) {
 
         it('should emit acceptance telemetry for README generation', async () => {
             await performAction('generate', getSessionStub)
-            await new Promise((resolve) => setTimeout(resolve, 100))
+            await waitForTelemetryCall(sendDocTelemetrySpy, 1)
+            await performAction('accept', getSessionStub)
             const expectedEvent = createExpectedEvent({
                 type: 'acceptance',
                 ...EventMetrics.INITIAL_README,
                 interactionType: 'GENERATE_README',
                 conversationId: conversationID,
             })
-
-            await performAction('accept', getSessionStub)
+            await waitForTelemetryCall(sendDocTelemetrySpy, 2)
             await assertTelemetry({
                 spy: sendDocTelemetrySpy,
                 expectedEvent,
@@ -231,7 +233,7 @@ for (let i = 0; i < 1000; i++) {
         })
         it('should emit generation telemetry for README update', async () => {
             await performAction('update', getSessionStub)
-
+            await waitForTelemetryCall(sendDocTelemetrySpy, 1)
             const expectedEvent = createExpectedEvent({
                 type: 'generation',
                 ...EventMetrics.REPO_STRUCTURE,
@@ -247,13 +249,13 @@ for (let i = 0; i < 1000; i++) {
         })
         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))
+            await waitForTelemetryCall(sendDocTelemetrySpy, 1)
 
             modifiedReadme = ReadmeBuilder.createReadmeWithDataFlow()
             await updateFilePaths(session, modifiedReadme, uploadID, docScheme, controllerSetup.workspaceFolder)
 
             await performAction('makeChanges', getSessionStub, 'add data flow section')
-
+            await waitForTelemetryCall(sendDocTelemetrySpy, 2)
             const expectedEvent = createExpectedEvent({
                 type: 'generation',
                 ...EventMetrics.DATA_FLOW,
@@ -272,16 +274,17 @@ for (let i = 0; i < 1000; i++) {
 
         it('should emit acceptance telemetry for README update', async () => {
             await performAction('update', getSessionStub)
-            await new Promise((resolve) => setTimeout(resolve, 100))
+            await waitForTelemetryCall(sendDocTelemetrySpy, 1)
 
+            await performAction('accept', getSessionStub)
+            await waitForTelemetryCall(sendDocTelemetrySpy, 2)
             const expectedEvent = createExpectedEvent({
                 type: 'acceptance',
                 ...EventMetrics.REPO_STRUCTURE,
                 interactionType: 'UPDATE_README',
                 conversationId: conversationID,
             })
 
-            await performAction('accept', getSessionStub)
             await assertTelemetry({
                 spy: sendDocTelemetrySpy,
                 expectedEvent,
@@ -292,7 +295,7 @@ for (let i = 0; i < 1000; i++) {
 
         it('should emit generation telemetry for README edit', async () => {
             await performAction('edit', getSessionStub, 'add repository structure section')
-
+            await waitForTelemetryCall(sendDocTelemetrySpy, 1)
             const expectedEvent = createExpectedEvent({
                 type: 'generation',
                 ...EventMetrics.REPO_STRUCTURE,
@@ -308,16 +311,16 @@ for (let i = 0; i < 1000; i++) {
         })
         it('should emit acceptance telemetry for README edit', async () => {
             await performAction('edit', getSessionStub, 'add repository structure section')
-            await new Promise((resolve) => setTimeout(resolve, 100))
-
+            await waitForTelemetryCall(sendDocTelemetrySpy, 1)
+            await performAction('accept', getSessionStub)
+            await waitForTelemetryCall(sendDocTelemetrySpy, 2)
             const expectedEvent = createExpectedEvent({
                 type: 'acceptance',
                 ...EventMetrics.REPO_STRUCTURE,
                 interactionType: 'EDIT_README',
                 conversationId: conversationID,
             })
 
-            await performAction('accept', getSessionStub)
             await assertTelemetry({
                 spy: sendDocTelemetrySpy,
                 expectedEvent,

packages/core/src/test/amazonqDoc/utils.ts

@@ -21,6 +21,7 @@ import { createSessionConfig } from '../../amazonq/commons/session/sessionConfig
 import { DocGenerationTask } from '../../amazonqDoc/controllers/docGenerationTask'
 import { DocV2GenerationEvent, DocV2AcceptanceEvent } from '../../amazonqFeatureDev/client/featuredevproxyclient'
 import { FollowUpTypes } from '../../amazonq/commons/types'
+import { waitUntil } from '../../shared'
 
 export function createMessenger(): DocMessenger {
     return new DocMessenger(
@@ -200,7 +201,6 @@ export async function assertTelemetry(params: {
     type: 'generation' | 'acceptance'
     callIndex?: number
 }) {
-    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)
 }
@@ -247,3 +247,7 @@ export const FollowUpSequences = {
     makeChanges: [FollowUpTypes.MakeChanges],
     acceptContent: [FollowUpTypes.AcceptChanges],
 }
+
+export async function waitForTelemetryCall(spy: sinon.SinonStub, expectedCallCount: number) {
+    await waitUntil(() => Promise.resolve(spy.callCount >= expectedCallCount), { timeout: 2000, interval: 100 })
+}