@@ -8,8 +8,11 @@ import assert from 'assert'
88import sinon from 'sinon'
99import { DocPrepareCodeGenState } from '../../../amazonqDoc'
1010import { createMockSessionStateAction } from '../../amazonq/utils'
11-
1211import { createTestContext , setupTestHooks } from '../../amazonq/session/testSetup'
12+ import * as filesModule from '../../../amazonq/util/files'
13+ import { CodeGenBase } from '../../../amazonq/session/sessionState'
14+ import { RunCommandLogFileName } from '../../../amazonq/session/sessionState'
15+ import fss from '../../../shared/fs/fs'
1316
1417describe ( 'sessionStateDoc' , ( ) => {
1518 const context = createTestContext ( )
@@ -26,4 +29,168 @@ describe('sessionStateDoc', () => {
2629 } )
2730 } )
2831 } )
32+
33+ describe ( 'CodeGenBase generateCode log file handling' , ( ) => {
34+ class TestCodeGen extends CodeGenBase {
35+ public generatedFiles : any [ ] = [ ]
36+ constructor ( config : any , tabID : string ) {
37+ super ( config , tabID )
38+ }
39+ protected handleProgress ( _messenger : any ) : void {
40+ // No-op for test.
41+ }
42+ protected getScheme ( ) : string {
43+ return 'file'
44+ }
45+ protected getTimeoutErrorCode ( ) : string {
46+ return 'test_timeout'
47+ }
48+ protected handleGenerationComplete ( _messenger : any , newFileInfo : any [ ] ) : void {
49+ this . generatedFiles = newFileInfo
50+ }
51+ protected handleError ( _messenger : any , _codegenResult : any ) : Error {
52+ throw new Error ( 'handleError called' )
53+ }
54+ }
55+
56+ let fakeProxyClient : any
57+ let testConfig : any
58+ let fsMock : any
59+ let messengerMock : any
60+ let telemetryMock : any
61+ let loggerMock : any
62+ let testAction : any
63+ let registerNewFilesStub : any
64+ let mkdirStub : any
65+ let writeFileStub : any
66+
67+ beforeEach ( ( ) => {
68+ fakeProxyClient = {
69+ getCodeGeneration : sinon . stub ( ) . resolves ( {
70+ codeGenerationStatus : { status : 'Complete' } ,
71+ codeGenerationRemainingIterationCount : 0 ,
72+ codeGenerationTotalIterationCount : 1 ,
73+ } ) ,
74+ exportResultArchive : sinon . stub ( ) ,
75+ }
76+
77+ testConfig = {
78+ conversationId : 'conv_test' ,
79+ uploadId : 'upload_test' ,
80+ workspaceRoots : [ '/workspace' ] ,
81+ proxyClient : fakeProxyClient ,
82+ }
83+
84+ fsMock = {
85+ writeFile : sinon . stub ( ) . resolves ( ) ,
86+ }
87+
88+ messengerMock = { sendAnswer : sinon . spy ( ) }
89+
90+ telemetryMock = {
91+ setCodeGenerationResult : sinon . spy ( ) ,
92+ setNumberOfFilesGenerated : sinon . spy ( ) ,
93+ setAmazonqNumberOfReferences : sinon . spy ( ) ,
94+ setGenerateCodeIteration : sinon . spy ( ) ,
95+ setGenerateCodeLastInvocationTime : sinon . spy ( ) ,
96+ recordUserCodeGenerationTelemetry : sinon . spy ( ) ,
97+ }
98+
99+ loggerMock = {
100+ addContent : sinon . spy ( ) ,
101+ }
102+
103+ testAction = {
104+ fs : fsMock ,
105+ messenger : messengerMock ,
106+ logger : loggerMock ,
107+ tokenSource : { token : { isCancellationRequested : false , onCancellationRequested : ( ) => { } } } ,
108+ }
109+
110+ registerNewFilesStub = sinon
111+ . stub ( filesModule , 'registerNewFiles' )
112+ . callsFake ( ( _ , newFileContents : any [ ] ) => {
113+ return newFileContents
114+ } )
115+
116+ mkdirStub = sinon . stub ( fss , 'mkdir' ) . resolves ( )
117+ writeFileStub = sinon . stub ( fss , 'writeFile' ) . resolves ( )
118+ } )
119+
120+ afterEach ( ( ) => {
121+ sinon . restore ( )
122+ } )
123+
124+ it ( 'adds the log content to logger if present and removes it from new files' , async ( ) => {
125+ const logFileInfo = {
126+ zipFilePath : RunCommandLogFileName ,
127+ fileContent : 'newLog' ,
128+ }
129+ const otherFile = { zipFilePath : 'other.ts' , fileContent : 'other content' }
130+ fakeProxyClient . exportResultArchive . resolves ( {
131+ newFileContents : [ logFileInfo , otherFile ] ,
132+ deletedFiles : [ ] ,
133+ references : [ ] ,
134+ } )
135+
136+ const testCodeGen = new TestCodeGen ( testConfig , 'tab1' )
137+
138+ const result = await testCodeGen . generateCode ( {
139+ messenger : messengerMock ,
140+ fs : fsMock ,
141+ codeGenerationId : 'codegen1' ,
142+ telemetry : telemetryMock ,
143+ workspaceFolders : [ ] as any ,
144+ action : testAction ,
145+ } )
146+
147+ sinon . assert . calledWith ( loggerMock . addContent , 'newLog' )
148+
149+ sinon . assert . notCalled ( mkdirStub )
150+ sinon . assert . notCalled ( writeFileStub )
151+
152+ sinon . assert . calledWith (
153+ registerNewFilesStub ,
154+ fsMock ,
155+ [ otherFile ] ,
156+ testConfig . uploadId ,
157+ [ ] as any ,
158+ testConfig . conversationId ,
159+ 'file'
160+ )
161+ assert . deepStrictEqual ( result . newFiles , [ otherFile ] )
162+ } )
163+
164+ it ( 'skips log file handling if log file is not present' , async ( ) => {
165+ const file1 = { zipFilePath : 'file1.ts' , fileContent : 'content1' }
166+ fakeProxyClient . exportResultArchive . resolves ( {
167+ newFileContents : [ file1 ] ,
168+ deletedFiles : [ ] ,
169+ references : [ ] ,
170+ } )
171+
172+ const testCodeGen = new TestCodeGen ( testConfig , 'tab1' )
173+
174+ const result = await testCodeGen . generateCode ( {
175+ messenger : messengerMock ,
176+ fs : fsMock ,
177+ codeGenerationId : 'codegen2' ,
178+ telemetry : telemetryMock ,
179+ workspaceFolders : [ ] as any ,
180+ action : testAction ,
181+ } )
182+
183+ sinon . assert . notCalled ( loggerMock . addContent )
184+ sinon . assert . calledWith (
185+ registerNewFilesStub ,
186+ fsMock ,
187+ [ file1 ] ,
188+ testConfig . uploadId ,
189+ [ ] as any ,
190+ testConfig . conversationId ,
191+ 'file'
192+ )
193+ assert . deepStrictEqual ( result . newFiles , [ file1 ] )
194+ } )
195+ } )
29196} )
0 commit comments