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