@@ -6,6 +6,7 @@ import assert from 'assert'
66import * as codewhispererClient from 'aws-core-vscode/codewhisperer'
77import * as EditorContext from 'aws-core-vscode/codewhisperer'
88import {
9+ createMockDocument ,
910 createMockTextEditor ,
1011 createMockClientRequest ,
1112 resetCodeWhispererGlobalVariables ,
@@ -15,6 +16,27 @@ import {
1516} from 'aws-core-vscode/test'
1617import { globals } from 'aws-core-vscode/shared'
1718import { GenerateCompletionsRequest } from 'aws-core-vscode/codewhisperer'
19+ import * as vscode from 'vscode'
20+
21+ export function createNotebookCell (
22+ document : vscode . TextDocument = createMockDocument ( 'def example():\n return "test"' ) ,
23+ kind : vscode . NotebookCellKind = vscode . NotebookCellKind . Code ,
24+ notebook : vscode . NotebookDocument = { } as any ,
25+ index : number = 0 ,
26+ outputs : vscode . NotebookCellOutput [ ] = [ ] ,
27+ metadata : { readonly [ key : string ] : any } = { } ,
28+ executionSummary ?: vscode . NotebookCellExecutionSummary
29+ ) : vscode . NotebookCell {
30+ return {
31+ document,
32+ kind,
33+ notebook,
34+ index,
35+ outputs,
36+ metadata,
37+ executionSummary,
38+ }
39+ }
1840
1941describe ( 'editorContext' , function ( ) {
2042 let telemetryEnabledDefault : boolean
@@ -63,6 +85,44 @@ describe('editorContext', function () {
6385 }
6486 assert . deepStrictEqual ( actual , expected )
6587 } )
88+
89+ it ( 'in a notebook, includes context from other cells' , async function ( ) {
90+ const cells : vscode . NotebookCellData [ ] = [
91+ new vscode . NotebookCellData ( vscode . NotebookCellKind . Markup , 'Previous cell' , 'python' ) ,
92+ new vscode . NotebookCellData (
93+ vscode . NotebookCellKind . Code ,
94+ 'import numpy as np\nimport pandas as pd\n\ndef analyze_data(df):\n # Current cell with cursor here' ,
95+ 'python'
96+ ) ,
97+ new vscode . NotebookCellData (
98+ vscode . NotebookCellKind . Code ,
99+ '# Process the data\nresult = analyze_data(df)\nprint(result)' ,
100+ 'python'
101+ ) ,
102+ ]
103+
104+ const document = await vscode . workspace . openNotebookDocument (
105+ 'jupyter-notebook' ,
106+ new vscode . NotebookData ( cells )
107+ )
108+ const editor : any = {
109+ document : document . cellAt ( 1 ) . document ,
110+ selection : { active : new vscode . Position ( 4 , 13 ) } ,
111+ }
112+
113+ const actual = EditorContext . extractContextForCodeWhisperer ( editor )
114+ const expected : codewhispererClient . FileContext = {
115+ filename : 'Untitled-1.py' ,
116+ programmingLanguage : {
117+ languageName : 'python' ,
118+ } ,
119+ leftFileContent :
120+ '# Previous cell\nimport numpy as np\nimport pandas as pd\n\ndef analyze_data(df):\n # Current' ,
121+ rightFileContent :
122+ ' cell with cursor here\n# Process the data\nresult = analyze_data(df)\nprint(result)\n' ,
123+ }
124+ assert . deepStrictEqual ( actual , expected )
125+ } )
66126 } )
67127
68128 describe ( 'getFileName' , function ( ) {
@@ -115,6 +175,165 @@ describe('editorContext', function () {
115175 } )
116176 } )
117177
178+ describe ( 'getNotebookCellContext' , function ( ) {
179+ it ( 'Should return cell text for python code cells when language is python' , function ( ) {
180+ const mockCodeCell = createNotebookCell ( createMockDocument ( 'def example():\n return "test"' ) )
181+ const result = EditorContext . getNotebookCellContext ( mockCodeCell , 'python' )
182+ assert . strictEqual ( result , 'def example():\n return "test"' )
183+ } )
184+
185+ it ( 'Should return java comments for python code cells when language is java' , function ( ) {
186+ const mockCodeCell = createNotebookCell ( createMockDocument ( 'def example():\n return "test"' ) )
187+ const result = EditorContext . getNotebookCellContext ( mockCodeCell , 'java' )
188+ assert . strictEqual ( result , '// def example():\n// return "test"' )
189+ } )
190+
191+ it ( 'Should return python comments for java code cells when language is python' , function ( ) {
192+ const mockCodeCell = createNotebookCell ( createMockDocument ( 'println(1 + 1);' , 'somefile.ipynb' , 'java' ) )
193+ const result = EditorContext . getNotebookCellContext ( mockCodeCell , 'python' )
194+ assert . strictEqual ( result , '# println(1 + 1);' )
195+ } )
196+
197+ it ( 'Should add python comment prefixes for markdown cells when language is python' , function ( ) {
198+ const mockMarkdownCell = createNotebookCell (
199+ createMockDocument ( '# Heading\nThis is a markdown cell' ) ,
200+ vscode . NotebookCellKind . Markup
201+ )
202+ const result = EditorContext . getNotebookCellContext ( mockMarkdownCell , 'python' )
203+ assert . strictEqual ( result , '# # Heading\n# This is a markdown cell' )
204+ } )
205+
206+ it ( 'Should add java comment prefixes for markdown cells when language is java' , function ( ) {
207+ const mockMarkdownCell = createNotebookCell (
208+ createMockDocument ( '# Heading\nThis is a markdown cell' ) ,
209+ vscode . NotebookCellKind . Markup
210+ )
211+ const result = EditorContext . getNotebookCellContext ( mockMarkdownCell , 'java' )
212+ assert . strictEqual ( result , '// # Heading\n// This is a markdown cell' )
213+ } )
214+ } )
215+
216+ describe ( 'getNotebookCellsSliceContext' , function ( ) {
217+ it ( 'Should extract content from cells in reverse order up to maxLength from prefix cells' , function ( ) {
218+ const mockCells = [
219+ createNotebookCell ( createMockDocument ( 'First cell content' ) ) ,
220+ createNotebookCell ( createMockDocument ( 'Second cell content' ) ) ,
221+ createNotebookCell ( createMockDocument ( 'Third cell content' ) ) ,
222+ ]
223+
224+ const result = EditorContext . getNotebookCellsSliceContext ( mockCells , 100 , 'python' , false )
225+ assert . strictEqual ( result , 'First cell content\nSecond cell content\nThird cell content\n' )
226+ } )
227+
228+ it ( 'Should extract content from cells in reverse order up to maxLength from suffix cells' , function ( ) {
229+ const mockCells = [
230+ createNotebookCell ( createMockDocument ( 'First cell content' ) ) ,
231+ createNotebookCell ( createMockDocument ( 'Second cell content' ) ) ,
232+ createNotebookCell ( createMockDocument ( 'Third cell content' ) ) ,
233+ ]
234+
235+ const result = EditorContext . getNotebookCellsSliceContext ( mockCells , 100 , 'python' , true )
236+ assert . strictEqual ( result , 'First cell content\nSecond cell content\nThird cell content\n' )
237+ } )
238+
239+ it ( 'Should respect maxLength parameter from prefix cells' , function ( ) {
240+ const mockCells = [
241+ createNotebookCell ( createMockDocument ( 'First' ) ) ,
242+ createNotebookCell ( createMockDocument ( 'Second' ) ) ,
243+ createNotebookCell ( createMockDocument ( 'Third' ) ) ,
244+ createNotebookCell ( createMockDocument ( 'Fourth' ) ) ,
245+ ]
246+ // Should only include part of second cell and the last two cells
247+ const result = EditorContext . getNotebookCellsSliceContext ( mockCells , 15 , 'python' , false )
248+ assert . strictEqual ( result , 'd\nThird\nFourth\n' )
249+ } )
250+
251+ it ( 'Should respect maxLength parameter from suffix cells' , function ( ) {
252+ const mockCells = [
253+ createNotebookCell ( createMockDocument ( 'First' ) ) ,
254+ createNotebookCell ( createMockDocument ( 'Second' ) ) ,
255+ createNotebookCell ( createMockDocument ( 'Third' ) ) ,
256+ createNotebookCell ( createMockDocument ( 'Fourth' ) ) ,
257+ ]
258+
259+ // Should only include first cell and part of second cell
260+ const result = EditorContext . getNotebookCellsSliceContext ( mockCells , 15 , 'python' , true )
261+ assert . strictEqual ( result , 'First\nSecond\nTh' )
262+ } )
263+
264+ it ( 'Should handle empty cells array from prefix cells' , function ( ) {
265+ const result = EditorContext . getNotebookCellsSliceContext ( [ ] , 100 , 'python' , false )
266+ assert . strictEqual ( result , '' )
267+ } )
268+
269+ it ( 'Should handle empty cells array from suffix cells' , function ( ) {
270+ const result = EditorContext . getNotebookCellsSliceContext ( [ ] , 100 , 'python' , true )
271+ assert . strictEqual ( result , '' )
272+ } )
273+
274+ it ( 'Should add python comments to markdown prefix cells' , function ( ) {
275+ const mockCells = [
276+ createNotebookCell ( createMockDocument ( '# Heading\nThis is markdown' ) , vscode . NotebookCellKind . Markup ) ,
277+ createNotebookCell ( createMockDocument ( 'def example():\n return "test"' ) ) ,
278+ ]
279+ const result = EditorContext . getNotebookCellsSliceContext ( mockCells , 100 , 'python' , false )
280+ assert . strictEqual ( result , '# # Heading\n# This is markdown\ndef example():\n return "test"\n' )
281+ } )
282+
283+ it ( 'Should add python comments to markdown suffix cells' , function ( ) {
284+ const mockCells = [
285+ createNotebookCell ( createMockDocument ( '# Heading\nThis is markdown' ) , vscode . NotebookCellKind . Markup ) ,
286+ createNotebookCell ( createMockDocument ( 'def example():\n return "test"' ) ) ,
287+ ]
288+
289+ const result = EditorContext . getNotebookCellsSliceContext ( mockCells , 100 , 'python' , true )
290+ assert . strictEqual ( result , '# # Heading\n# This is markdown\ndef example():\n return "test"\n' )
291+ } )
292+
293+ it ( 'Should add java comments to markdown and python prefix cells when language is java' , function ( ) {
294+ const mockCells = [
295+ createNotebookCell ( createMockDocument ( '# Heading\nThis is markdown' ) , vscode . NotebookCellKind . Markup ) ,
296+ createNotebookCell ( createMockDocument ( 'def example():\n return "test"' ) ) ,
297+ ]
298+ const result = EditorContext . getNotebookCellsSliceContext ( mockCells , 100 , 'java' , false )
299+ assert . strictEqual ( result , '// # Heading\n// This is markdown\n// def example():\n// return "test"\n' )
300+ } )
301+
302+ it ( 'Should add java comments to markdown and python suffix cells when language is java' , function ( ) {
303+ const mockCells = [
304+ createNotebookCell ( createMockDocument ( '# Heading\nThis is markdown' ) , vscode . NotebookCellKind . Markup ) ,
305+ createNotebookCell ( createMockDocument ( 'println(1 + 1);' , 'somefile.ipynb' , 'java' ) ) ,
306+ ]
307+
308+ const result = EditorContext . getNotebookCellsSliceContext ( mockCells , 100 , 'java' , true )
309+ assert . strictEqual ( result , '// # Heading\n// This is markdown\nprintln(1 + 1);\n' )
310+ } )
311+
312+ it ( 'Should handle code prefix cells with different languages' , function ( ) {
313+ const mockCells = [
314+ createNotebookCell (
315+ createMockDocument ( 'println(1 + 1);' , 'somefile.ipynb' , 'java' ) ,
316+ vscode . NotebookCellKind . Code
317+ ) ,
318+ createNotebookCell ( createMockDocument ( 'def example():\n return "test"' ) ) ,
319+ ]
320+ const result = EditorContext . getNotebookCellsSliceContext ( mockCells , 100 , 'python' , false )
321+ assert . strictEqual ( result , '# println(1 + 1);\ndef example():\n return "test"\n' )
322+ } )
323+
324+ it ( 'Should handle code suffix cells with different languages' , function ( ) {
325+ const mockCells = [
326+ createNotebookCell (
327+ createMockDocument ( 'println(1 + 1);' , 'somefile.ipynb' , 'java' ) ,
328+ vscode . NotebookCellKind . Code
329+ ) ,
330+ createNotebookCell ( createMockDocument ( 'def example():\n return "test"' ) ) ,
331+ ]
332+ const result = EditorContext . getNotebookCellsSliceContext ( mockCells , 100 , 'python' , true )
333+ assert . strictEqual ( result , '# println(1 + 1);\ndef example():\n return "test"\n' )
334+ } )
335+ } )
336+
118337 describe ( 'validateRequest' , function ( ) {
119338 it ( 'Should return false if request filename.length is invalid' , function ( ) {
120339 const req = createMockClientRequest ( )
0 commit comments