@@ -21,6 +21,7 @@ import {
2121import { IDisposable } from '../../platform/common/types' ;
2222import * as notebookUpdater from '../../kernels/execution/notebookUpdater' ;
2323import { createMockedNotebookDocument } from '../../test/datascience/editor-integration/helpers' ;
24+ import { WrappedError } from '../../platform/errors/types' ;
2425import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes' ;
2526
2627suite ( 'DeepnoteNotebookCommandListener' , ( ) => {
@@ -976,5 +977,191 @@ suite('DeepnoteNotebookCommandListener', () => {
976977 ) ;
977978 } ) ;
978979 } ) ;
980+
981+ suite ( 'addChartBlock' , ( ) => {
982+ test ( 'should add chart block at the end when no selection exists' , async ( ) => {
983+ // Setup mocks
984+ const { editor, document } = createMockEditor ( [ ] , undefined ) ;
985+ const { chainStub, executeCommandStub, getCapturedNotebookEdits } =
986+ mockNotebookUpdateAndExecute ( editor ) ;
987+
988+ // Call the method
989+ await commandListener . addChartBlock ( ) ;
990+
991+ const capturedNotebookEdits = getCapturedNotebookEdits ( ) ;
992+
993+ // Verify chainWithPendingUpdates was called
994+ assert . isTrue ( chainStub . calledOnce , 'chainWithPendingUpdates should be called once' ) ;
995+ assert . equal ( chainStub . firstCall . args [ 0 ] , document , 'Should be called with correct document' ) ;
996+
997+ // Verify the edits were captured
998+ assert . isNotNull ( capturedNotebookEdits , 'Notebook edits should be captured' ) ;
999+ assert . isDefined ( capturedNotebookEdits , 'Notebook edits should be defined' ) ;
1000+
1001+ const editsArray = capturedNotebookEdits ! ;
1002+ assert . equal ( editsArray . length , 1 , 'Should have one notebook edit' ) ;
1003+
1004+ const notebookEdit = editsArray [ 0 ] as any ;
1005+ assert . equal ( notebookEdit . newCells . length , 1 , 'Should insert one cell' ) ;
1006+
1007+ const newCell = notebookEdit . newCells [ 0 ] ;
1008+ assert . equal ( newCell . kind , NotebookCellKind . Code , 'Should be a code cell' ) ;
1009+ assert . equal ( newCell . languageId , 'json' , 'Should have json language' ) ;
1010+
1011+ // Verify cell content is valid JSON with correct structure
1012+ const content = JSON . parse ( newCell . value ) ;
1013+ assert . equal ( content . variable , 'df_1' , 'Should have correct variable name' ) ;
1014+ assert . property ( content , 'spec' , 'Should have spec property' ) ;
1015+ assert . property ( content , 'filters' , 'Should have filters property' ) ;
1016+
1017+ // Verify the spec has the correct Vega-Lite structure
1018+ assert . equal ( content . spec . mark , 'line' , 'Should be a line chart' ) ;
1019+ assert . equal (
1020+ content . spec . $schema ,
1021+ 'https://vega.github.io/schema/vega-lite/v5.json' ,
1022+ 'Should have Vega-Lite schema'
1023+ ) ;
1024+ assert . deepStrictEqual ( content . spec . data , { values : [ ] } , 'Should have empty data array' ) ;
1025+ assert . property ( content . spec , 'encoding' , 'Should have encoding property' ) ;
1026+ assert . property ( content . spec . encoding , 'x' , 'Should have x encoding' ) ;
1027+ assert . property ( content . spec . encoding , 'y' , 'Should have y encoding' ) ;
1028+
1029+ // Verify metadata structure
1030+ assert . property ( newCell . metadata , '__deepnotePocket' , 'Should have __deepnotePocket metadata' ) ;
1031+ assert . equal ( newCell . metadata . __deepnotePocket . type , 'visualization' , 'Should have visualization type' ) ;
1032+
1033+ // Verify reveal and selection were set
1034+ assert . isTrue ( ( editor . revealRange as sinon . SinonStub ) . calledOnce , 'Should reveal the new cell range' ) ;
1035+ const revealCall = ( editor . revealRange as sinon . SinonStub ) . firstCall ;
1036+ assert . equal ( revealCall . args [ 0 ] . start , 0 , 'Should reveal correct range start' ) ;
1037+ assert . equal ( revealCall . args [ 0 ] . end , 1 , 'Should reveal correct range end' ) ;
1038+ assert . equal ( revealCall . args [ 1 ] , 0 , 'Should use NotebookEditorRevealType.Default (value 0)' ) ;
1039+
1040+ // Verify notebook.cell.edit command was executed
1041+ assert . isTrue (
1042+ executeCommandStub . calledWith ( 'notebook.cell.edit' ) ,
1043+ 'Should execute notebook.cell.edit command'
1044+ ) ;
1045+ } ) ;
1046+
1047+ test ( 'should add chart block after selection when selection exists' , async ( ) => {
1048+ // Setup mocks
1049+ const existingCells = [ createMockCell ( '{}' ) , createMockCell ( '{}' ) ] ;
1050+ const selection = new NotebookRange ( 0 , 1 ) ;
1051+ const { editor } = createMockEditor ( existingCells , selection ) ;
1052+ const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute ( editor ) ;
1053+
1054+ // Call the method
1055+ await commandListener . addChartBlock ( ) ;
1056+
1057+ const capturedNotebookEdits = getCapturedNotebookEdits ( ) ;
1058+
1059+ // Verify chainWithPendingUpdates was called
1060+ assert . isTrue ( chainStub . calledOnce , 'chainWithPendingUpdates should be called once' ) ;
1061+
1062+ // Verify a cell was inserted
1063+ assert . isNotNull ( capturedNotebookEdits , 'Notebook edits should be captured' ) ;
1064+ const notebookEdit = capturedNotebookEdits ! [ 0 ] as any ;
1065+ assert . equal ( notebookEdit . newCells . length , 1 , 'Should insert one cell' ) ;
1066+ assert . equal ( notebookEdit . newCells [ 0 ] . languageId , 'json' , 'Should be JSON cell' ) ;
1067+ } ) ;
1068+
1069+ test ( 'should use hardcoded variable name df_1' , async ( ) => {
1070+ // Setup mocks with existing df variables
1071+ const existingCells = [
1072+ createMockCell ( '{ "deepnote_variable_name": "df_1" }' ) ,
1073+ createMockCell ( '{ "variable": "df_2" }' )
1074+ ] ;
1075+ const { editor } = createMockEditor ( existingCells , undefined ) ;
1076+ const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute ( editor ) ;
1077+
1078+ // Call the method
1079+ await commandListener . addChartBlock ( ) ;
1080+
1081+ const capturedNotebookEdits = getCapturedNotebookEdits ( ) ;
1082+ const notebookEdit = capturedNotebookEdits ! [ 0 ] as any ;
1083+ const newCell = notebookEdit . newCells [ 0 ] ;
1084+
1085+ // Verify variable name is always df_1
1086+ const content = JSON . parse ( newCell . value ) ;
1087+ assert . equal ( content . variable , 'df_1' , 'Should always use df_1' ) ;
1088+ } ) ;
1089+
1090+ test ( 'should always use df_1 regardless of existing variables' , async ( ) => {
1091+ // Setup mocks with various existing variables
1092+ const existingCells = [
1093+ createMockCell ( '{ "deepnote_variable_name": "input_10" }' ) ,
1094+ createMockCell ( '{ "deepnote_variable_name": "df_5" }' ) ,
1095+ createMockCell ( '{ "variable": "df_2" }' )
1096+ ] ;
1097+ const { editor } = createMockEditor ( existingCells , undefined ) ;
1098+ const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute ( editor ) ;
1099+
1100+ // Call the method
1101+ await commandListener . addChartBlock ( ) ;
1102+
1103+ const capturedNotebookEdits = getCapturedNotebookEdits ( ) ;
1104+ const notebookEdit = capturedNotebookEdits ! [ 0 ] as any ;
1105+ const newCell = notebookEdit . newCells [ 0 ] ;
1106+
1107+ // Verify variable name is always df_1
1108+ const content = JSON . parse ( newCell . value ) ;
1109+ assert . equal ( content . variable , 'df_1' , 'Should always use df_1' ) ;
1110+ } ) ;
1111+
1112+ test ( 'should insert at correct position in the middle of notebook' , async ( ) => {
1113+ // Setup mocks
1114+ const existingCells = [ createMockCell ( '{}' ) , createMockCell ( '{}' ) , createMockCell ( '{}' ) ] ;
1115+ const selection = new NotebookRange ( 1 , 2 ) ;
1116+ const { editor } = createMockEditor ( existingCells , selection ) ;
1117+ const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute ( editor ) ;
1118+
1119+ // Call the method
1120+ await commandListener . addChartBlock ( ) ;
1121+
1122+ const capturedNotebookEdits = getCapturedNotebookEdits ( ) ;
1123+
1124+ // Verify chainWithPendingUpdates was called
1125+ assert . isTrue ( chainStub . calledOnce , 'chainWithPendingUpdates should be called once' ) ;
1126+
1127+ // Verify a cell was inserted
1128+ assert . isNotNull ( capturedNotebookEdits , 'Notebook edits should be captured' ) ;
1129+ const notebookEdit = capturedNotebookEdits ! [ 0 ] as any ;
1130+ assert . equal ( notebookEdit . newCells . length , 1 , 'Should insert one cell' ) ;
1131+ assert . equal ( notebookEdit . newCells [ 0 ] . languageId , 'json' , 'Should be JSON cell' ) ;
1132+ } ) ;
1133+
1134+ test ( 'should throw error when no active editor exists' , async ( ) => {
1135+ // Setup: no active editor
1136+ Object . defineProperty ( window , 'activeNotebookEditor' , {
1137+ value : undefined ,
1138+ configurable : true ,
1139+ writable : true
1140+ } ) ;
1141+
1142+ // Call the method and expect rejection
1143+ await assert . isRejected (
1144+ commandListener . addChartBlock ( ) ,
1145+ WrappedError ,
1146+ 'No active notebook editor found'
1147+ ) ;
1148+ } ) ;
1149+
1150+ test ( 'should throw error when chainWithPendingUpdates fails' , async ( ) => {
1151+ // Setup mocks
1152+ const { editor } = createMockEditor ( [ ] , undefined ) ;
1153+ Object . defineProperty ( window , 'activeNotebookEditor' , {
1154+ value : editor ,
1155+ configurable : true ,
1156+ writable : true
1157+ } ) ;
1158+
1159+ // Mock chainWithPendingUpdates to return false
1160+ sandbox . stub ( notebookUpdater , 'chainWithPendingUpdates' ) . resolves ( false ) ;
1161+
1162+ // Call the method and expect rejection
1163+ await assert . isRejected ( commandListener . addChartBlock ( ) , WrappedError , 'Failed to insert chart block' ) ;
1164+ } ) ;
1165+ } ) ;
9791166 } ) ;
9801167} ) ;
0 commit comments