notebook variable data source tests (#203753)

* indexed children tests

* one more assert

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');
+ });
+});