@@ -8,8 +8,12 @@ import assert from 'assert'
8
8
import sinon from 'sinon'
9
9
import { DocPrepareCodeGenState } from '../../../amazonqDoc'
10
10
import { createMockSessionStateAction } from '../../amazonq/utils'
11
-
12
11
import { 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'
13
17
14
18
describe ( 'sessionStateDoc' , ( ) => {
15
19
const context = createTestContext ( )
@@ -26,4 +30,167 @@ describe('sessionStateDoc', () => {
26
30
} )
27
31
} )
28
32
} )
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
+ } )
29
196
} )
0 commit comments