Add test for notebook result grid streaming (#12539)

* start testing convertAllData

* add test for convertAllData method

* clean up code

src/sql/workbench/contrib/notebook/test/browser/dataResourceDataProvider.test.ts

@@ -13,7 +13,7 @@ import * as uuid from 'uuid';
 import * as sinon from 'sinon';
 import { DataResourceDataProvider } from '../../browser/outputs/gridOutput.component';
 import { IDataResource } from 'sql/workbench/services/notebook/browser/sql/sqlSessionManager';
-import { ResultSetSummary } from 'sql/workbench/services/query/common/query';
+import { ResultSetSubset, ResultSetSummary } from 'sql/workbench/services/query/common/query';
 import { TestNotificationService } from 'vs/platform/notification/test/common/testNotificationService';
 import { TestFileDialogService, TestEditorService } from 'vs/workbench/test/browser/workbenchTestServices';
 import { TestContextService } from 'vs/workbench/test/common/workbenchTestServices';
@@ -23,6 +23,7 @@ import { InstantiationService } from 'vs/platform/instantiation/common/instantia
 import { URI } from 'vs/base/common/uri';
 import { CellModel } from 'sql/workbench/services/notebook/browser/models/cell';
 import { createandLoadNotebookModel } from 'sql/workbench/contrib/notebook/test/browser/cellToolbarActions.test';
+import QueryRunner from 'sql/workbench/services/query/common/queryRunner';
 
 export class TestSerializationProvider implements azdata.SerializationProvider {
 	providerId: string;
@@ -47,67 +48,69 @@ export class TestSerializationProvider implements azdata.SerializationProvider {
 }
 
 suite('Data Resource Data Provider', function () {
-	let tempFolderPath: string;
 	let fileDialogService: TypeMoq.Mock<TestFileDialogService>;
 	let serializer: ResultSerializer;
-	let dataResourceDataProvider: DataResourceDataProvider;
+	let notificationService: TestNotificationService;
+	let serializationService: SerializationService;
+	let instantiationService: TypeMoq.Mock<InstantiationService>;
+	let cellModel = TypeMoq.Mock.ofType(CellModel);
+
+	// Create test data with two rows and two columns
+	let source: IDataResource = {
+		data: [{ 0: '1', 1: '2' }, { 0: '3', 1: '4' }],
+		schema: { fields: [{ name: 'col1' }, { name: 'col2' }] }
+	};
+	let resultSet: ResultSetSummary = {
+		batchId: 0,
+		columnInfo: [{ columnName: 'col1' }, { columnName: 'col2' }],
+		complete: true,
+		id: 0,
+		rowCount: 2
+	};
 
 	suiteSetup(async () => {
-		// Create test data with two rows and two columns
-		let source: IDataResource = {
-			data: [{ 0: '1', 1: '2' }, { 0: '3', 1: '4' }],
-			schema: { fields: [{ name: 'col1' }, { name: 'col2' }] }
-		};
-		let resultSet: ResultSetSummary = {
-			batchId: 0,
-			columnInfo: [{ columnName: 'col1' }, { columnName: 'col2' }],
-			complete: true,
-			id: 0,
-			rowCount: 2
-		};
-		let cellModel = TypeMoq.Mock.ofType(CellModel);
 		let notebookModel = await createandLoadNotebookModel();
 		cellModel.setup(x => x.notebookModel).returns(() => notebookModel);
-		tempFolderPath = path.join(os.tmpdir(), `TestDataResourceDataProvider_${uuid.v4()}`);
-		await fs.mkdir(tempFolderPath);
 
 		// Mock services
 		let editorService = TypeMoq.Mock.ofType(TestEditorService, TypeMoq.MockBehavior.Strict);
 		editorService.setup(x => x.openEditor(TypeMoq.It.isAny())).returns(() => Promise.resolve(undefined));
 		let contextService = new TestContextService();
 		fileDialogService = TypeMoq.Mock.ofType(TestFileDialogService, TypeMoq.MockBehavior.Strict);
-		let _notificationService = new TestNotificationService();
+		notificationService = new TestNotificationService();
-		let _serializationService = new SerializationService(undefined, undefined); //_connectionService _capabilitiesService
+		serializationService = new SerializationService(undefined, undefined); //_connectionService _capabilitiesService
-		_serializationService.registerProvider('testProviderId', new TestSerializationProvider());
+		serializationService.registerProvider('testProviderId', new TestSerializationProvider());
 		serializer = new ResultSerializer(
 			undefined, // IQueryManagementService
 			undefined, // IConfigurationService
 			editorService.object,
 			contextService,
 			fileDialogService.object,
-			_notificationService,
+			notificationService,
 			undefined // IOpenerService
 		);
-		let _instantiationService = TypeMoq.Mock.ofType(InstantiationService, TypeMoq.MockBehavior.Strict);
+		instantiationService = TypeMoq.Mock.ofType(InstantiationService, TypeMoq.MockBehavior.Strict);
-		_instantiationService.setup(x => x.createInstance(TypeMoq.It.isValue(ResultSerializer)))
+		instantiationService.setup(x => x.createInstance(TypeMoq.It.isValue(ResultSerializer)))
 			.returns(() => serializer);
-		dataResourceDataProvider = new DataResourceDataProvider(
+	});
+
+	test('serializeResults call is successful', async function (): Promise<void> {
+		let tempFolderPath = path.join(os.tmpdir(), `TestDataResourceDataProvider_${uuid.v4()}`);
+		await fs.mkdir(tempFolderPath);
+		let dataResourceDataProvider = new DataResourceDataProvider(
 			0, // batchId
 			0, // id
 			undefined, // QueryRunner
 			source,
 			resultSet,
 			cellModel.object,
-			_notificationService,
+			notificationService,
 			undefined, // IClipboardService
 			undefined, // IConfigurationService
 			undefined, // ITextResourcePropertiesService
-			_serializationService,
+			serializationService,
-			_instantiationService.object
+			instantiationService.object
 		);
-	});
-
-	test('serializeResults call is successful', async function (): Promise<void> {
 		let noHeadersFile = URI.file(path.join(tempFolderPath, 'result_noHeaders.csv'));
 		let fileDialogServiceStub = sinon.stub(fileDialogService.object, 'showSaveDialog').returns(Promise.resolve(noHeadersFile));
 		let serializerStub = sinon.stub(serializer, 'getBasicSaveParameters').returns({ resultFormat: SaveFormat.CSV as string, includeHeaders: false });
@@ -128,5 +131,38 @@ suite('Data Resource Data Provider', function () {
 		const withHeadersResult = await fs.readFile(withHeadersFile.fsPath);
 		assert.equal(withHeadersResult.toString(), 'col1 col2 \n1 2 \n3 4 \n', 'result data should include headers');
 	});
-});
 
+	test('convertAllData correctly converts row data to mimetype and html', async function (): Promise<void> {
+		let resultSetSubset: ResultSetSubset = {
+			rowCount: 2,
+			rows: [[{ displayValue: '1' }, { displayValue: '2' }], [{ displayValue: '3' }, { displayValue: '4' }]]
+		};
+		let queryRunner: TypeMoq.Mock<QueryRunner> = TypeMoq.Mock.ofType(QueryRunner);
+		queryRunner.setup(x => x.getQueryRows(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(resultSetSubset));
+		let dataResourceDataProvider = new DataResourceDataProvider(
+			0, // batchId
+			0, // id
+			queryRunner.object,
+			source,
+			resultSet,
+			cellModel.object,
+			notificationService,
+			undefined, // IClipboardService
+			undefined, // IConfigurationService
+			undefined, // ITextResourcePropertiesService
+			serializationService,
+			instantiationService.object
+		);
+		let spy = sinon.spy(cellModel.object, 'updateOutputData');
+		let expectedData = {
+			'application/vnd.dataresource+json': {
+				data: [{ 0: '1', 1: '2' }, { 0: '3', 1: '4' }],
+				schema: { fields: [{ name: 'col1' }, { name: 'col2' }] }
+			},
+			'text/html': ['<table>', '<tr><th>col1</th><th>col2</th></tr>', '<tr><td>1</td><td>2</td></tr>', '<tr><td>3</td><td>4</td></tr>', '</table>']
+		};
+		await dataResourceDataProvider.convertAllData(resultSet);
+		sinon.assert.calledOnce(spy);
+		sinon.assert.calledWithExactly(spy, 0, 0, expectedData);
+	});
+});