@@ -8,8 +8,12 @@ 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'
16+ import path from 'path'
1317
1418describe ( 'sessionStateDoc' , ( ) => {
1519 const context = createTestContext ( )
@@ -26,4 +30,167 @@ describe('sessionStateDoc', () => {
2630 } )
2731 } )
2832 } )
33+
34+ describe ( 'CodeGenBase generateCode log file handling' , ( ) => {
35+ // Minimal implementation of TestCodeGen to test generateCode.
36+ class TestCodeGen extends CodeGenBase {
37+ public generatedFiles : any [ ] = [ ]
38+ constructor ( config : any , tabID : string ) {
39+ super ( config , tabID )
40+ }
41+ protected handleProgress ( _messenger : any ) : void {
42+ // No-op for test.
43+ }
44+ protected getScheme ( ) : string {
45+ return 'file'
46+ }
47+ protected getTimeoutErrorCode ( ) : string {
48+ return 'test_timeout'
49+ }
50+ protected handleGenerationComplete ( _messenger : any , newFileInfo : any [ ] ) : void {
51+ this . generatedFiles = newFileInfo
52+ }
53+ protected handleError ( _messenger : any , _codegenResult : any ) : Error {
54+ throw new Error ( 'handleError called' )
55+ }
56+ }
57+
58+ let fakeProxyClient : any
59+ let testConfig : any
60+ let fsMock : any
61+ let messengerMock : any
62+ let telemetryMock : any
63+ let testAction : any
64+ let registerNewFilesStub : any
65+ let mkdirStub : any
66+ let writeFileStub : any
67+
68+ beforeEach ( ( ) => {
69+ fakeProxyClient = {
70+ getCodeGeneration : sinon . stub ( ) . resolves ( {
71+ codeGenerationStatus : { status : 'Complete' } ,
72+ codeGenerationRemainingIterationCount : 0 ,
73+ codeGenerationTotalIterationCount : 1 ,
74+ } ) ,
75+ exportResultArchive : sinon . stub ( ) ,
76+ }
77+
78+ testConfig = {
79+ conversationId : 'conv_test' ,
80+ uploadId : 'upload_test' ,
81+ workspaceRoots : [ '/workspace' ] ,
82+ proxyClient : fakeProxyClient ,
83+ }
84+
85+ fsMock = {
86+ writeFile : sinon . stub ( ) . resolves ( ) ,
87+ }
88+
89+ messengerMock = { sendAnswer : sinon . spy ( ) }
90+
91+ telemetryMock = {
92+ setCodeGenerationResult : sinon . spy ( ) ,
93+ setNumberOfFilesGenerated : sinon . spy ( ) ,
94+ setAmazonqNumberOfReferences : sinon . spy ( ) ,
95+ setGenerateCodeIteration : sinon . spy ( ) ,
96+ setGenerateCodeLastInvocationTime : sinon . spy ( ) ,
97+ recordUserCodeGenerationTelemetry : sinon . spy ( ) ,
98+ }
99+
100+ testAction = {
101+ fs : fsMock ,
102+ messenger : messengerMock ,
103+ tokenSource : { token : { isCancellationRequested : false , onCancellationRequested : ( ) => { } } } ,
104+ }
105+
106+ registerNewFilesStub = sinon
107+ . stub ( filesModule , 'registerNewFiles' )
108+ . callsFake ( ( _ , newFileContents : any [ ] ) => {
109+ return newFileContents
110+ } )
111+
112+ mkdirStub = sinon . stub ( fss , 'mkdir' ) . resolves ( )
113+ writeFileStub = sinon . stub ( fss , 'writeFile' ) . resolves ( )
114+ } )
115+
116+ afterEach ( ( ) => {
117+ sinon . restore ( )
118+ } )
119+
120+ it ( 'writes the log file if present and removes it from new files' , async ( ) => {
121+ const logFileInfo = {
122+ zipFilePath : RunCommandLogFileName ,
123+ fileContent : 'newLog' ,
124+ }
125+ const otherFile = { zipFilePath : 'other.ts' , fileContent : 'other content' }
126+ fakeProxyClient . exportResultArchive . resolves ( {
127+ newFileContents : [ logFileInfo , otherFile ] ,
128+ deletedFiles : [ ] ,
129+ references : [ ] ,
130+ } )
131+
132+ const testCodeGen = new TestCodeGen ( testConfig , 'tab1' )
133+
134+ const result = await testCodeGen . generateCode ( {
135+ messenger : messengerMock ,
136+ fs : fsMock ,
137+ codeGenerationId : 'codegen1' ,
138+ telemetry : telemetryMock ,
139+ workspaceFolders : [ ] as any ,
140+ action : testAction ,
141+ } )
142+
143+ const expectedFilePath = path . join ( testConfig . workspaceRoots [ 0 ] , RunCommandLogFileName )
144+ const fileUri = vscode . Uri . file ( expectedFilePath )
145+
146+ // Verify mkdir and writeFile have been called with the expected arguments.
147+ sinon . assert . calledWith ( mkdirStub , path . dirname ( fileUri . fsPath ) )
148+ sinon . assert . calledWith ( writeFileStub , fileUri . fsPath , 'newLog' )
149+
150+ sinon . assert . calledWith (
151+ registerNewFilesStub ,
152+ fsMock ,
153+ [ otherFile ] ,
154+ testConfig . uploadId ,
155+ [ ] as any ,
156+ testConfig . conversationId ,
157+ 'file'
158+ )
159+ assert . deepStrictEqual ( result . newFiles , [ otherFile ] )
160+ } )
161+
162+ it ( 'skips log file handling if log file is not present' , async ( ) => {
163+ const file1 = { zipFilePath : 'file1.ts' , fileContent : 'content1' }
164+ fakeProxyClient . exportResultArchive . resolves ( {
165+ newFileContents : [ file1 ] ,
166+ deletedFiles : [ ] ,
167+ references : [ ] ,
168+ } )
169+
170+ const testCodeGen = new TestCodeGen ( testConfig , 'tab1' )
171+
172+ const result = await testCodeGen . generateCode ( {
173+ messenger : messengerMock ,
174+ fs : fsMock ,
175+ codeGenerationId : 'codegen2' ,
176+ telemetry : telemetryMock ,
177+ workspaceFolders : [ ] as any ,
178+ action : testAction ,
179+ } )
180+
181+ // Verify that mkdir and writeFile were not called because no log file exists.
182+ sinon . assert . notCalled ( mkdirStub )
183+ sinon . assert . notCalled ( writeFileStub )
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