notebook variable data source tests (microsoft#203753)

* indexed children tests

* one more assert

Author: amunger

src/vs/workbench/contrib/notebook/test/browser/notebookVariablesDataSource.test.ts

@@ -0,0 +1,116 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+import * as assert from 'assert';
+import { AsyncIterableObject, AsyncIterableSource } from 'vs/base/common/async';
+import { CancellationToken } from 'vs/base/common/cancellation';
+import { URI } from 'vs/base/common/uri';
+import { mock } from 'vs/base/test/common/mock';
+import { ensureNoDisposablesAreLeakedInTestSuite } from 'vs/base/test/common/utils';
+import { INotebookVariableElement, NotebookVariableDataSource } from 'vs/workbench/contrib/notebook/browser/contrib/notebookVariables/notebookVariablesDataSource';
+import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
+import { INotebookKernel, INotebookKernelService, VariablesResult } from 'vs/workbench/contrib/notebook/common/notebookKernelService';
+
+
+suite('NotebookVariableDataSource', () => {
+	let dataSource: NotebookVariableDataSource;
+	const notebookModel = { uri: 'one.ipynb', languages: ['python'] } as unknown as NotebookTextModel;
+	let provideVariablesCalled = false;
+
+	type VariablesResultWithAction = VariablesResult & { action?: () => void };
+	let results: VariablesResultWithAction[] = [
+		{ id: 1, name: 'a', value: '1', hasNamedChildren: false, indexedChildrenCount: 0 },
+	];
+
+	const kernel = new class extends mock<INotebookKernel>() {
+		override hasVariableProvider = true;
+		override provideVariables(
+			notebookUri: URI,
+			parentId: number | undefined,
+			kind: 'named' | 'indexed',
+			start: number,
+			token: CancellationToken
+		): AsyncIterableObject<VariablesResult> {
+			provideVariablesCalled = true;
+			const source = new AsyncIterableSource<VariablesResult>();
+			for (let i = 0; i < results.length; i++) {
+				if (token.isCancellationRequested) {
+					break;
+				}
+				if (results[i].action) {
+					results[i].action!();
+				}
+				source.emitOne(results[i]);
+			}
+
+			setTimeout(() => source.resolve(), 0);
+			return source.asyncIterable;
+		}
+	};
+
+	const kernelService = new class extends mock<INotebookKernelService>() {
+		override getMatchingKernel(notebook: NotebookTextModel) {
+			return { selected: kernel, all: [], suggestions: [], hidden: [] };
+		}
+	};
+
+	ensureNoDisposablesAreLeakedInTestSuite();
+
+	setup(() => {
+		provideVariablesCalled = false;
+		dataSource = new NotebookVariableDataSource(kernelService);
+	});
+
+	test('Root element should return children', async () => {
+		const variables = await dataSource.getChildren({ kind: 'root', notebook: notebookModel });
+
+		assert.strictEqual(variables.length, 1);
+	});
+
+	test('Get children of list element', async () => {
+		const parent = { kind: 'variable', notebook: notebookModel, id: '1', extHostId: 1, name: 'list', value: '[...]', hasNamedChildren: false, indexedChildrenCount: 5 } as INotebookVariableElement;
+		results = [
+			{ id: 2, name: 'first', value: '1', hasNamedChildren: false, indexedChildrenCount: 0 },
+			{ id: 3, name: 'second', value: '2', hasNamedChildren: false, indexedChildrenCount: 0 },
+			{ id: 4, name: 'third', value: '3', hasNamedChildren: false, indexedChildrenCount: 0 },
+			{ id: 5, name: 'fourth', value: '4', hasNamedChildren: false, indexedChildrenCount: 0 },
+			{ id: 6, name: 'fifth', value: '5', hasNamedChildren: false, indexedChildrenCount: 0 },
+		];
+
+		const variables = await dataSource.getChildren(parent);
+
+		assert.strictEqual(variables.length, 5);
+	});
+
+	test('Get children for large list', async () => {
+		const parent = { kind: 'variable', notebook: notebookModel, id: '1', extHostId: 1, name: 'list', value: '[...]', hasNamedChildren: false, indexedChildrenCount: 2000 } as INotebookVariableElement;
+		results = [];
+
+		const variables = await dataSource.getChildren(parent);
+
+		assert(variables.length > 1, 'We should have results for groups of children');
+		assert(!provideVariablesCalled, 'provideVariables should not be called');
+		assert.equal(variables[0].extHostId, parent.extHostId, 'ExtHostId should match the parent since we will use it to get the real children');
+	});
+
+	test('Cancel while enumerating through children', async () => {
+		const parent = { kind: 'variable', notebook: notebookModel, id: '1', extHostId: 1, name: 'list', value: '[...]', hasNamedChildren: false, indexedChildrenCount: 10 } as INotebookVariableElement;
+		results = [
+			{ id: 2, name: 'first', value: '1', hasNamedChildren: false, indexedChildrenCount: 0 },
+			{ id: 3, name: 'second', value: '2', hasNamedChildren: false, indexedChildrenCount: 0 },
+			{ id: 4, name: 'third', value: '3', hasNamedChildren: false, indexedChildrenCount: 0 },
+			{ id: 5, name: 'fourth', value: '4', hasNamedChildren: false, indexedChildrenCount: 0 },
+			{ id: 5, name: 'fifth', value: '4', hasNamedChildren: false, indexedChildrenCount: 0, action: () => dataSource.cancel() } as VariablesResult,
+			{ id: 7, name: 'sixth', value: '6', hasNamedChildren: false, indexedChildrenCount: 0 },
+			{ id: 8, name: 'seventh', value: '7', hasNamedChildren: false, indexedChildrenCount: 0 },
+			{ id: 9, name: 'eighth', value: '8', hasNamedChildren: false, indexedChildrenCount: 0 },
+			{ id: 10, name: 'ninth', value: '9', hasNamedChildren: false, indexedChildrenCount: 0 },
+			{ id: 11, name: 'tenth', value: '10', hasNamedChildren: false, indexedChildrenCount: 0 },
+		];
+
+		const variables = await dataSource.getChildren(parent);
+
+		assert.equal(variables.length, 5, 'Iterating should have been cancelled');
+	});
+});