@@ -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,173 @@ 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+ workspaceFolders : { } as any ,
83+ proxyClient : fakeProxyClient ,
84+ currentCodeGenerationId : 'test_gen' ,
85+ }
86+
87+ fsMock = {
88+ writeFile : sinon . stub ( ) . resolves ( ) ,
89+ }
90+
91+ messengerMock = { sendAnswer : sinon . spy ( ) }
92+
93+ telemetryMock = {
94+ setCodeGenerationResult : sinon . spy ( ) ,
95+ setNumberOfFilesGenerated : sinon . spy ( ) ,
96+ setAmazonqNumberOfReferences : sinon . spy ( ) ,
97+ setGenerateCodeIteration : sinon . spy ( ) ,
98+ setGenerateCodeLastInvocationTime : sinon . spy ( ) ,
99+ recordUserCodeGenerationTelemetry : sinon . spy ( ) ,
100+ }
101+
102+ testAction = {
103+ fs : fsMock ,
104+ messenger : messengerMock ,
105+ tokenSource : { token : { isCancellationRequested : false , onCancellationRequested : ( ) => { } } } ,
106+ telemetry : telemetryMock ,
107+ uploadHistory : { } ,
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 ( 'writes the log file 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+ const expectedFilePath = path . join ( testConfig . workspaceRoots [ 0 ] , RunCommandLogFileName )
148+ const fileUri = vscode . Uri . file ( expectedFilePath )
149+
150+ // Verify mkdir and writeFile have been called with the expected arguments.
151+ sinon . assert . calledWith ( mkdirStub , path . dirname ( fileUri . fsPath ) )
152+ sinon . assert . calledWith ( writeFileStub , fileUri . fsPath , 'newLog' )
153+
154+ // Verify that registerNewFiles received only the "otherFile" (i.e. log file removed)
155+ sinon . assert . calledWith (
156+ registerNewFilesStub ,
157+ fsMock ,
158+ [ otherFile ] ,
159+ testConfig . uploadId ,
160+ { } as any ,
161+ testConfig . conversationId ,
162+ 'file'
163+ )
164+ assert . deepStrictEqual ( result . newFiles , [ otherFile ] )
165+ } )
166+
167+ it ( 'skips log file handling if log file is not present' , async ( ) => {
168+ const file1 = { zipFilePath : 'file1.ts' , fileContent : 'content1' }
169+ fakeProxyClient . exportResultArchive . resolves ( {
170+ newFileContents : [ file1 ] ,
171+ deletedFiles : [ ] ,
172+ references : [ ] ,
173+ } )
174+
175+ const testCodeGen = new TestCodeGen ( testConfig , 'tab1' )
176+
177+ const result = await testCodeGen . generateCode ( {
178+ messenger : messengerMock ,
179+ fs : fsMock ,
180+ codeGenerationId : 'codegen2' ,
181+ telemetry : telemetryMock ,
182+ workspaceFolders : { } as any ,
183+ action : testAction ,
184+ } )
185+
186+ // Verify that mkdir and writeFile were not called because no log file exists.
187+ sinon . assert . notCalled ( mkdirStub )
188+ sinon . assert . notCalled ( writeFileStub )
189+ // Verify that the registered new files are the same as provided.
190+ sinon . assert . calledWith (
191+ registerNewFilesStub ,
192+ fsMock ,
193+ [ file1 ] ,
194+ testConfig . uploadId ,
195+ { } as any ,
196+ testConfig . conversationId ,
197+ 'file'
198+ )
199+ assert . deepStrictEqual ( result . newFiles , [ file1 ] )
200+ } )
201+ } )
29202} )
0 commit comments