@@ -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,173 @@ 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
+ 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
+ } )
29
202
} )
0 commit comments