Adding tests for DataResourceDataProvider (#11743)

* tests for getRow method

* use suiteSetup

* added more tests

* revert changes to gridOutput.component.ts

* mocked out services for DataResourceDataProvider

* add test for serializeResults

* change path to fsPath

* fix param type and use fs promises

* use async in startSerialization

* mock ContextService

src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts

@@ -225,7 +225,7 @@ class DataResourceTable extends GridTableBase<any> {
 	}
 }
 
-class DataResourceDataProvider implements IGridDataProvider {
+export class DataResourceDataProvider implements IGridDataProvider {
 	private rows: ICellValue[][];
 	constructor(source: IDataResource,
 		private resultSet: ResultSetSummary,

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

@@ -0,0 +1,127 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as assert from 'assert';
+import * as TypeMoq from 'typemoq';
+import * as azdata from 'azdata';
+import * as os from 'os';
+import * as path from 'path';
+import * as fs from 'fs-extra';
+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 { 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';
+import { SerializationService } from 'sql/platform/serialization/common/serializationService';
+import { SaveFormat, ResultSerializer } from 'sql/workbench/services/query/common/resultSerializer';
+import { InstantiationService } from 'vs/platform/instantiation/common/instantiationService';
+import { URI } from 'vs/base/common/uri';
+
+export class TestSerializationProvider implements azdata.SerializationProvider {
+	providerId: string;
+	constructor(providerId: string = 'providerId') { }
+
+	// Write data to file
+	async startSerialization(requestParams: azdata.SerializeDataStartRequestParams): Promise<azdata.SerializeDataResult> {
+		let data: string = '';
+		requestParams.rows.forEach((row) => {
+			row.forEach((element) => {
+				data += element.displayValue + ' ';
+			});
+			data += '\n';
+		});
+		await fs.promises.writeFile(requestParams.filePath, data);
+		return Promise.resolve({ succeeded: true, messages: undefined });
+	}
+
+	continueSerialization(requestParams: azdata.SerializeDataContinueRequestParams): Thenable<azdata.SerializeDataResult> {
+		return Promise.resolve(undefined);
+	}
+}
+
+suite('Data Resource Data Provider', function () {
+	let tempFolderPath: string;
+	let fileDialogService: TypeMoq.Mock<TestFileDialogService>;
+	let serializer: ResultSerializer;
+	let dataResourceDataProvider: DataResourceDataProvider;
+
+	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 documentUri = 'untitled:Notebook-0';
+		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();
+		let _serializationService = new SerializationService(undefined, undefined); //_connectionService _capabilitiesService
+		_serializationService.registerProvider('testProviderId', new TestSerializationProvider());
+		serializer = new ResultSerializer(
+			undefined, // IQueryManagementService
+			undefined, // IConfigurationService
+			editorService.object,
+			contextService,
+			fileDialogService.object,
+			_notificationService,
+			undefined // IOpenerService
+		);
+		let _instantiationService = TypeMoq.Mock.ofType(InstantiationService, TypeMoq.MockBehavior.Strict);
+		_instantiationService.setup(x => x.createInstance(TypeMoq.It.isValue(ResultSerializer)))
+			.returns(() => serializer);
+
+		dataResourceDataProvider = new DataResourceDataProvider(
+			source,
+			resultSet,
+			documentUri,
+			_notificationService,
+			undefined, // IClipboardService
+			undefined, // IConfigurationService
+			undefined, // ITextResourcePropertiesService
+			_serializationService,
+			_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 });
+		await dataResourceDataProvider.serializeResults(SaveFormat.CSV, undefined);
+		fileDialogServiceStub.restore();
+		serializerStub.restore();
+
+		let withHeadersFile = URI.file(path.join(tempFolderPath, 'result_withHeaders.csv'));
+		fileDialogServiceStub = sinon.stub(fileDialogService.object, 'showSaveDialog').returns(Promise.resolve(withHeadersFile));
+		serializerStub = sinon.stub(serializer, 'getBasicSaveParameters').returns({ resultFormat: SaveFormat.CSV as string, includeHeaders: true });
+		await dataResourceDataProvider.serializeResults(SaveFormat.CSV, undefined);
+		fileDialogServiceStub.restore();
+		serializerStub.restore();
+
+		const noHeadersResult = await fs.readFile(noHeadersFile.fsPath);
+		assert.equal(noHeadersResult.toString(), '1 2 \n3 4 \n', 'result data should not include headers');
+
+		const withHeadersResult = await fs.readFile(withHeadersFile.fsPath);
+		assert.equal(withHeadersResult.toString(), 'col1 col2 \n1 2 \n3 4 \n', 'result data should include headers');
+	});
+});
+