@@ -10,6 +10,7 @@ import { DocPrepareCodeGenState } from '../../../amazonqDoc'
1010import { createMockSessionStateAction } from '../../amazonq/utils'
1111
1212import { createTestContext , setupTestHooks } from '../../amazonq/session/testSetup'
13+ const filesModule = require ( '../../../amazonq/session/sessionState' )
1314
1415describe ( 'sessionStateDoc' , ( ) => {
1516 const context = createTestContext ( )
@@ -26,4 +27,198 @@ describe('sessionStateDoc', () => {
2627 } )
2728 } )
2829 } )
30+
31+ describe ( 'CodeGenBase generateCode log file handling' , ( ) => {
32+ const RunCommandLogFileName = '.amazonq/dev/run_command.log'
33+
34+ // Minimal stub for registerNewFiles to return newFileContents without the log file entry
35+ const registerNewFilesStub = sinon . stub ( )
36+ registerNewFilesStub . callsFake ( ( fs : any , newFileContents : any [ ] ) => {
37+ // In our test scenario, assume registerNewFiles returns the files as is (log file already spliced out)
38+ return newFileContents
39+ } )
40+
41+ // Fake CodeGenBase subclass for testing
42+ class TestCodeGen extends ( require ( '../../../amazonq/session/sessionState' ) as any ) . CodeGenBase {
43+ public generatedFiles : any [ ] = [ ]
44+ constructor ( config : any , tabID : string ) {
45+ super ( config , tabID )
46+ }
47+ protected handleProgress ( messenger : any , action : any , detail ?: string ) : void {
48+ // no-op
49+ }
50+ protected getScheme ( ) : string {
51+ return 'file'
52+ }
53+ protected getTimeoutErrorCode ( ) : string {
54+ return 'test_timeout'
55+ }
56+ protected handleGenerationComplete ( messenger : any , newFileInfo : any [ ] , action : any ) : void {
57+ this . generatedFiles = newFileInfo
58+ }
59+ protected handleError ( messenger : any , codegenResult : any ) : Error {
60+ throw new Error ( 'handleError called' )
61+ }
62+ }
63+
64+ // Common objects for testing
65+ let testConfig : any
66+ let fakeProxyClient : any
67+ let fsMock : any
68+ let telemetryMock : any
69+ let messengerMock : any
70+ let testAction : any
71+
72+ beforeEach ( ( ) => {
73+ fakeProxyClient = {
74+ getCodeGeneration : sinon . stub ( ) . resolves ( {
75+ codeGenerationStatus : { status : 'Complete' } ,
76+ codeGenerationRemainingIterationCount : 0 ,
77+ codeGenerationTotalIterationCount : 1 ,
78+ } ) ,
79+ exportResultArchive : sinon . stub ( ) ,
80+ }
81+
82+ testConfig = {
83+ conversationId : 'conv1' ,
84+ uploadId : 'upload1' ,
85+ workspaceRoots : [ '/workspace' ] ,
86+ proxyClient : fakeProxyClient ,
87+ }
88+
89+ fsMock = {
90+ stat : sinon . stub ( ) ,
91+ readFile : sinon . stub ( ) ,
92+ writeFile : sinon . stub ( ) ,
93+ }
94+
95+ telemetryMock = {
96+ setCodeGenerationResult : sinon . spy ( ) ,
97+ setNumberOfFilesGenerated : sinon . spy ( ) ,
98+ setAmazonqNumberOfReferences : sinon . spy ( ) ,
99+ setGenerateCodeIteration : sinon . spy ( ) ,
100+ setGenerateCodeLastInvocationTime : sinon . spy ( ) ,
101+ recordUserCodeGenerationTelemetry : sinon . spy ( ) ,
102+ }
103+
104+ messengerMock = {
105+ sendAnswer : sinon . spy ( ) ,
106+ }
107+
108+ testAction = {
109+ telemetry : telemetryMock ,
110+ fs : fsMock ,
111+ messenger : messengerMock ,
112+ uploadHistory : { } ,
113+ tokenSource : { token : { isCancellationRequested : false , onCancellationRequested : ( ) => { } } } ,
114+ }
115+ } )
116+
117+ afterEach ( ( ) => {
118+ sinon . restore ( )
119+ } )
120+
121+ it ( 'should append existing log file content to new log content and remove the log file entry when file exists' , async ( ) => {
122+ // Prepare exportResultArchive to return a log file and another file
123+ const logFileInfo = {
124+ name : RunCommandLogFileName ,
125+ content : 'newLog' ,
126+ }
127+ const otherFile = { name : 'other.ts' , content : 'other' }
128+
129+ fakeProxyClient . exportResultArchive . resolves ( {
130+ newFileContents : [ logFileInfo , otherFile ] ,
131+ deletedFiles : [ ] ,
132+ references : [ ] ,
133+ } )
134+
135+ // Stub fs.stat to simulate that file exists
136+ fsMock . stat . resolves ( { } )
137+ // Stub fs.readFile to simulate existing file content
138+ fsMock . readFile . resolves ( Buffer . from ( 'existing' ) )
139+
140+ // Stub registerNewFiles function used inside generateCode
141+ // const registerNewFilesOriginal = filesModule.registerNewFiles
142+ filesModule . registerNewFiles = registerNewFilesStub
143+
144+ const testCodeGen = new TestCodeGen ( testConfig , 'tab1' )
145+ // Force a single polling iteration by simulating COMPLETE
146+
147+ // const result =
148+ await testCodeGen . generateCode ( {
149+ messenger : messengerMock ,
150+ fs : fsMock ,
151+ codeGenerationId : 'codegen1' ,
152+ telemetry : telemetryMock ,
153+ workspaceFolders : { } ,
154+ action : testAction ,
155+ } )
156+
157+ // Check that fs.stat and fs.readFile were called with the expected fileUri
158+ const expectedFilePath = `${ testConfig . workspaceRoots [ 0 ] } /${ RunCommandLogFileName } `
159+ const fileUri = vscode . Uri . file ( expectedFilePath )
160+ sinon . assert . calledWith ( fsMock . stat , fileUri )
161+ sinon . assert . calledWith ( fsMock . readFile , fileUri )
162+
163+ // Verify that the log file content was updated correctly: "existing" + "newLog"
164+ // And that the log file info was spliced out so that registerNewFiles gets only the other file.
165+ assert . strictEqual ( logFileInfo . content , 'existingnewLog' )
166+ assert . deepStrictEqual ( testCodeGen . generatedFiles , [ otherFile ] )
167+
168+ // Restore registerNewFiles to its original implementation
169+ sinon . restore ( )
170+ } )
171+
172+ it ( 'should create a new log file when log file does not exist' , async ( ) => {
173+ // Prepare exportResultArchive to return a log file and another file
174+ const logFileInfo = {
175+ name : RunCommandLogFileName ,
176+ content : 'newLog' ,
177+ }
178+ const otherFile = { name : 'other.ts' , content : 'other' }
179+
180+ fakeProxyClient . exportResultArchive . resolves ( {
181+ newFileContents : [ logFileInfo , otherFile ] ,
182+ deletedFiles : [ ] ,
183+ references : [ ] ,
184+ } )
185+
186+ // Stub fs.stat to simulate that file does not exist
187+ fsMock . stat . rejects ( new Error ( 'Not found' ) )
188+ // fs.writeFile should be called in this case
189+ fsMock . writeFile . resolves ( )
190+
191+ // Stub registerNewFiles function
192+ // sinon.replace(filesModule, 'registerNewFiles', registerNewFilesStub)
193+ filesModule . registerNewFiles = registerNewFilesStub
194+
195+ const testCodeGen = new TestCodeGen ( testConfig , 'tab1' )
196+
197+ // const result =
198+ await testCodeGen . generateCode ( {
199+ messenger : messengerMock ,
200+ fs : fsMock ,
201+ codeGenerationId : 'codegen2' ,
202+ telemetry : telemetryMock ,
203+ workspaceFolders : { } ,
204+ action : testAction ,
205+ } )
206+
207+ // Verify that fs.stat was called and then fs.writeFile was invoked
208+ const expectedFilePath = `${ testConfig . workspaceRoots [ 0 ] } /${ RunCommandLogFileName } `
209+ const fileUri = vscode . Uri . file ( expectedFilePath )
210+ sinon . assert . calledWith ( fsMock . stat , fileUri )
211+ sinon . assert . calledWith ( fsMock . writeFile , fileUri , new TextEncoder ( ) . encode ( 'newLog' ) , {
212+ create : true ,
213+ overwrite : true ,
214+ } )
215+ // fs.readFile should not be called
216+ sinon . assert . notCalled ( fsMock . readFile )
217+
218+ // Verify the log file has been removed from new file contents
219+ assert . deepStrictEqual ( testCodeGen . generatedFiles , [ otherFile ] )
220+
221+ sinon . restore ( )
222+ } )
223+ } )
29224} )
0 commit comments