@@ -9,7 +9,13 @@ import * as path from 'path'
99import sinon from 'sinon'
1010import { waitUntil } from '../../../../shared/utilities/timeoutUtils'
1111import { ControllerSetup , createController , createSession , generateVirtualMemoryUri } from '../../utils'
12- import { CurrentWsFolders , DeletedFileInfo , NewFileInfo } from '../../../../amazonqFeatureDev/types'
12+ import {
13+ CurrentWsFolders ,
14+ DeletedFileInfo ,
15+ MetricDataOperationName ,
16+ MetricDataResult ,
17+ NewFileInfo ,
18+ } from '../../../../amazonqFeatureDev/types'
1319import { Session } from '../../../../amazonqFeatureDev/session/session'
1420import { Prompter } from '../../../../shared/ui/prompter'
1521import { assertTelemetry , toFile } from '../../../testUtil'
@@ -36,6 +42,7 @@ import { AuthUtil } from '../../../../codewhisperer'
3642import { featureDevScheme , featureName , messageWithConversationId } from '../../../../amazonqFeatureDev'
3743import { i18n } from '../../../../shared/i18n-helper'
3844import { FollowUpTypes } from '../../../../amazonq/commons/types'
45+ import { ToolkitError } from '../../../../shared'
3946
4047let mockGetCodeGeneration : sinon . SinonStub
4148describe ( 'Controller' , ( ) => {
@@ -395,7 +402,47 @@ describe('Controller', () => {
395402 } )
396403
397404 describe ( 'processUserChatMessage' , function ( ) {
398- async function fireChatMessage ( ) {
405+ // TODO: fix disablePreviousFileList error
406+ const runs = [
407+ { name : 'ContentLengthError' , error : new ContentLengthError ( ) } ,
408+ {
409+ name : 'MonthlyConversationLimitError' ,
410+ error : new MonthlyConversationLimitError ( 'Service Quota Exceeded' ) ,
411+ } ,
412+ {
413+ name : 'FeatureDevServiceErrorGuardrailsException' ,
414+ error : new FeatureDevServiceError (
415+ i18n ( 'AWS.amazonq.featureDev.error.codeGen.default' ) ,
416+ 'GuardrailsException'
417+ ) ,
418+ } ,
419+ {
420+ name : 'FeatureDevServiceErrorEmptyPatchException' ,
421+ error : new FeatureDevServiceError (
422+ i18n ( 'AWS.amazonq.featureDev.error.throttling' ) ,
423+ 'EmptyPatchException'
424+ ) ,
425+ } ,
426+ {
427+ name : 'FeatureDevServiceErrorThrottlingException' ,
428+ error : new FeatureDevServiceError (
429+ i18n ( 'AWS.amazonq.featureDev.error.codeGen.default' ) ,
430+ 'ThrottlingException'
431+ ) ,
432+ } ,
433+ { name : 'UploadCodeError' , error : new UploadCodeError ( '403: Forbiden' ) } ,
434+ { name : 'UserMessageNotFoundError' , error : new UserMessageNotFoundError ( ) } ,
435+ { name : 'TabIdNotFoundError' , error : new TabIdNotFoundError ( ) } ,
436+ { name : 'PrepareRepoFailedError' , error : new PrepareRepoFailedError ( ) } ,
437+ { name : 'PromptRefusalException' , error : new PromptRefusalException ( ) } ,
438+ { name : 'ZipFileError' , error : new ZipFileError ( ) } ,
439+ { name : 'CodeIterationLimitError' , error : new CodeIterationLimitError ( ) } ,
440+ { name : 'UploadURLExpired' , error : new UploadURLExpired ( ) } ,
441+ { name : 'NoChangeRequiredException' , error : new NoChangeRequiredException ( ) } ,
442+ { name : 'default' , error : new ToolkitError ( 'Default' , { code : 'Default' } ) } ,
443+ ]
444+
445+ async function fireChatMessage ( session : Session ) {
399446 const getSessionStub = sinon . stub ( controllerSetup . sessionStorage , 'getSession' ) . resolves ( session )
400447
401448 controllerSetup . emitters . processHumanChatMessage . fire ( {
@@ -415,6 +462,109 @@ describe('Controller', () => {
415462 } , { } )
416463 }
417464
465+ describe ( 'onCodeGeneration' , function ( ) {
466+ let session : any
467+ let sendMetricDataTelemetrySpy : sinon . SinonStub
468+
469+ const errorResultMapping = new Map ( [
470+ [ 'EmptyPatchException' , MetricDataResult . LlmFailure ] ,
471+ [ PromptRefusalException . name , MetricDataResult . Error ] ,
472+ [ NoChangeRequiredException . name , MetricDataResult . Error ] ,
473+ ] )
474+
475+ function getMetricResult ( error : ToolkitError ) : MetricDataResult {
476+ if ( error instanceof FeatureDevServiceError && error . code ) {
477+ return errorResultMapping . get ( error . code ) ?? MetricDataResult . Error
478+ }
479+ return errorResultMapping . get ( error . constructor . name ) ?? MetricDataResult . Fault
480+ }
481+
482+ async function createCodeGenState ( ) {
483+ mockGetCodeGeneration = sinon . stub ( ) . resolves ( { codeGenerationStatus : { status : 'Complete' } } )
484+
485+ const workspaceFolders = [ controllerSetup . workspaceFolder ] as CurrentWsFolders
486+ const testConfig = {
487+ conversationId : conversationID ,
488+ proxyClient : {
489+ createConversation : ( ) => sinon . stub ( ) ,
490+ createUploadUrl : ( ) => sinon . stub ( ) ,
491+ generatePlan : ( ) => sinon . stub ( ) ,
492+ startCodeGeneration : ( ) => sinon . stub ( ) ,
493+ getCodeGeneration : ( ) => mockGetCodeGeneration ( ) ,
494+ exportResultArchive : ( ) => sinon . stub ( ) ,
495+ } as unknown as FeatureDevClient ,
496+ workspaceRoots : [ '' ] ,
497+ uploadId : uploadID ,
498+ workspaceFolders,
499+ }
500+
501+ const codeGenState = new CodeGenState ( testConfig , getFilePaths ( controllerSetup ) , [ ] , [ ] , tabID , 0 , { } )
502+ const newSession = await createSession ( {
503+ messenger : controllerSetup . messenger ,
504+ sessionState : codeGenState ,
505+ conversationID,
506+ tabID,
507+ uploadID,
508+ scheme : featureDevScheme ,
509+ } )
510+ return newSession
511+ }
512+
513+ async function verifyException ( error : ToolkitError ) {
514+ sinon . stub ( session , 'send' ) . throws ( error )
515+
516+ await fireChatMessage ( session )
517+ await verifyMetricsCalled ( )
518+ assert . ok (
519+ sendMetricDataTelemetrySpy . calledWith (
520+ MetricDataOperationName . StartCodeGeneration ,
521+ MetricDataResult . Success
522+ )
523+ )
524+ const metricResult = getMetricResult ( error )
525+ assert . ok (
526+ sendMetricDataTelemetrySpy . calledWith ( MetricDataOperationName . EndCodeGeneration , metricResult )
527+ )
528+ }
529+
530+ async function verifyMetricsCalled ( ) {
531+ await waitUntil ( ( ) => Promise . resolve ( sendMetricDataTelemetrySpy . callCount >= 2 ) , { } )
532+ }
533+
534+ beforeEach ( async ( ) => {
535+ session = await createCodeGenState ( )
536+ sinon . stub ( session , 'preloader' ) . resolves ( )
537+ sendMetricDataTelemetrySpy = sinon . stub ( session , 'sendMetricDataTelemetry' )
538+ } )
539+
540+ it ( 'sends success operation telemetry' , async ( ) => {
541+ sinon . stub ( session , 'send' ) . resolves ( )
542+ sinon . stub ( session , 'sendLinesOfCodeGeneratedTelemetry' ) . resolves ( ) // Avoid sending extra telemetry
543+
544+ await fireChatMessage ( session )
545+ await verifyMetricsCalled ( )
546+
547+ assert . ok (
548+ sendMetricDataTelemetrySpy . calledWith (
549+ MetricDataOperationName . StartCodeGeneration ,
550+ MetricDataResult . Success
551+ )
552+ )
553+ assert . ok (
554+ sendMetricDataTelemetrySpy . calledWith (
555+ MetricDataOperationName . EndCodeGeneration ,
556+ MetricDataResult . Success
557+ )
558+ )
559+ } )
560+
561+ runs . forEach ( ( { name, error } ) => {
562+ it ( `sends failure operation telemetry on ${ name } ` , async ( ) => {
563+ await verifyException ( error )
564+ } )
565+ } )
566+ } )
567+
418568 describe ( 'processErrorChatMessage' , function ( ) {
419569 // TODO: fix disablePreviousFileList error
420570 const runs = [
@@ -446,12 +596,12 @@ describe('Controller', () => {
446596 return createUserFacingErrorMessage ( `${ featureName } request failed: ${ message } ` )
447597 }
448598
449- async function verifyException ( error : Error ) {
599+ async function verifyException ( error : ToolkitError ) {
450600 sinon . stub ( session , 'preloader' ) . throws ( error )
451601 const sendAnswerSpy = sinon . stub ( controllerSetup . messenger , 'sendAnswer' )
452602 const sendErrorMessageSpy = sinon . stub ( controllerSetup . messenger , 'sendErrorMessage' )
453603 const sendMonthlyLimitErrorSpy = sinon . stub ( controllerSetup . messenger , 'sendMonthlyLimitError' )
454- await fireChatMessage ( )
604+ await fireChatMessage ( session )
455605
456606 switch ( error . constructor . name ) {
457607 case ContentLengthError . name :
0 commit comments