Skip to content

Commit b106711

Browse files
committed
Fixes
1 parent 27b8e03 commit b106711

File tree

2 files changed

+82
-21
lines changed

2 files changed

+82
-21
lines changed

src/test/datascience/notebook/helper.ts

Lines changed: 71 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,8 @@ import {
4343
env,
4444
languages,
4545
window,
46-
workspace
46+
workspace,
47+
type CancellationToken
4748
} from 'vscode';
4849
import { DebugProtocol } from 'vscode-debugprotocol';
4950
import {
@@ -78,11 +79,11 @@ import {
7879
defaultNotebookFormat,
7980
isWebExtension
8081
} from '../../../platform/common/constants';
81-
import { dispose } from '../../../platform/common/utils/lifecycle';
82+
import { dispose, type DisposableStore } from '../../../platform/common/utils/lifecycle';
8283
import { getDisplayPath } from '../../../platform/common/platform/fs-paths';
8384
import { IFileSystem, IPlatformService } from '../../../platform/common/platform/types';
8485
import { GLOBAL_MEMENTO, IDisposable, IMemento } from '../../../platform/common/types';
85-
import { createDeferred, sleep } from '../../../platform/common/utils/async';
86+
import { createDeferred, raceTimeoutError, sleep } from '../../../platform/common/utils/async';
8687
import { DataScience } from '../../../platform/common/utils/localize';
8788
import { isWeb } from '../../../platform/common/utils/misc';
8889
import { openAndShowNotebook } from '../../../platform/common/utils/notebooks';
@@ -104,6 +105,7 @@ import {
104105
NotebookCellExecutionState,
105106
notebookCellExecutions
106107
} from '../../../platform/notebooks/cellExecutionStateService';
108+
import { raceCancellationError } from '../../../platform/common/cancellation';
107109

108110
// Running in Conda environments, things can be a little slower.
109111
export const defaultNotebookTestTimeout = 60_000;
@@ -982,6 +984,40 @@ export async function waitForExecutionCompletedSuccessfully(
982984
await sleep(100);
983985
}
984986

987+
export async function waitForExecutionCompletedSuccessfullyV2(
988+
cell: NotebookCell,
989+
disposables: IDisposable[] | DisposableStore
990+
) {
991+
const checkCompletedSuccessfully = () => {
992+
return cell.executionSummary?.success &&
993+
cell.executionSummary.executionOrder &&
994+
cell.executionSummary.timing?.endTime
995+
? true
996+
: false;
997+
};
998+
if (checkCompletedSuccessfully()) {
999+
return;
1000+
}
1001+
await new Promise((resolve) => {
1002+
const disposable = workspace.onDidChangeNotebookDocument((e) => {
1003+
if (e.notebook !== cell.notebook) {
1004+
return;
1005+
}
1006+
e.cellChanges.forEach(() => {
1007+
if (checkCompletedSuccessfully()) {
1008+
disposable.dispose();
1009+
resolve;
1010+
}
1011+
});
1012+
});
1013+
if (Array.isArray(disposables)) {
1014+
disposables.push(disposable);
1015+
} else {
1016+
disposables.add(disposable);
1017+
}
1018+
});
1019+
}
1020+
9851021
export async function waitForCompletions(
9861022
completionProvider: CompletionItemProvider,
9871023
cell: NotebookCell,
@@ -1220,6 +1256,38 @@ export async function waitForTextOutput(
12201256
.join(',\n')}`
12211257
);
12221258
}
1259+
export async function waitForTextOutputV2(
1260+
cell: NotebookCell,
1261+
text: string,
1262+
index: number = 0,
1263+
isExactMatch = true,
1264+
disposables: IDisposable[] | DisposableStore
1265+
) {
1266+
try {
1267+
assertHasTextOutputInVSCode(cell, text, index, isExactMatch);
1268+
return;
1269+
} catch {
1270+
//
1271+
}
1272+
await new Promise<void>((resolve) => {
1273+
const disposable = workspace.onDidChangeNotebookDocument((e) => {
1274+
if (e.notebook !== cell.notebook) {
1275+
return;
1276+
}
1277+
try {
1278+
assertHasTextOutputInVSCode(cell, text, index, isExactMatch);
1279+
resolve();
1280+
} catch {
1281+
//
1282+
}
1283+
});
1284+
if (Array.isArray(disposables)) {
1285+
disposables.push(disposable);
1286+
} else {
1287+
disposables.add(disposable);
1288+
}
1289+
});
1290+
}
12231291
export function assertNotHasTextOutputInVSCode(cell: NotebookCell, text: string, index: number, isExactMatch = true) {
12241292
const cellOutputs = cell.outputs;
12251293
assert.ok(cellOutputs, 'No output');

src/test/smoke/datascience.smoke.test.ts

Lines changed: 11 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ import { closeActiveWindows, initialize, initializeTest } from '../initialize.no
1111
import { captureScreenShot } from '../common';
1212
import { getCachedEnvironments } from '../../platform/interpreter/helpers';
1313
import { PythonExtension, type EnvironmentPath } from '@vscode/python-extension';
14+
import { waitForExecutionCompletedSuccessfullyV2, waitForTextOutputV2 } from '../datascience/notebook/helper';
15+
import { DisposableStore } from '../../platform/common/utils/lifecycle';
1416

1517
type JupyterApi = {
1618
openNotebook(uri: vscode.Uri, env: EnvironmentPath): Promise<void>;
@@ -19,6 +21,7 @@ type JupyterApi = {
1921
const timeoutForCellToRun = 3 * 60 * 1_000;
2022
suite('Smoke Tests', function () {
2123
this.timeout(timeoutForCellToRun);
24+
const disposableStore = new DisposableStore();
2225
suiteSetup(async function () {
2326
this.timeout(timeoutForCellToRun);
2427
if (!IS_SMOKE_TEST()) {
@@ -33,6 +36,7 @@ suite('Smoke Tests', function () {
3336
});
3437
suiteTeardown(closeActiveWindows);
3538
teardown(async function () {
39+
disposableStore.clear();
3640
traceInfo(`End Test ${this.currentTest?.title}`);
3741
if (this.currentTest?.isFailed()) {
3842
await captureScreenShot(this);
@@ -69,15 +73,15 @@ suite('Smoke Tests', function () {
6973
// }).timeout(timeoutForCellToRun);
7074

7175
test('Run Cell in Notebook', async function () {
72-
const cell = new vscode.NotebookCellData(vscode.NotebookCellKind.Code, 'print("Hello World")', 'python');
7376
const jupyterExt = vscode.extensions.getExtension<JupyterApi>(JVSC_EXTENSION_ID_FOR_TESTS);
7477
if (!jupyterExt) {
7578
throw new Error('Jupyter extension not found');
7679
}
80+
const cellData = new vscode.NotebookCellData(vscode.NotebookCellKind.Code, 'print("Hello World")', 'python');
7781
const [pythonEnv, { notebook }] = await Promise.all([
7882
PythonExtension.api().then((api) => api.environments.resolveEnvironment(PYTHON_PATH)),
7983
vscode.workspace
80-
.openNotebookDocument('jupyter-notebook', new vscode.NotebookData([cell]))
84+
.openNotebookDocument('jupyter-notebook', new vscode.NotebookData([cellData]))
8185
.then((notebook) => vscode.window.showNotebookDocument(notebook)),
8286
jupyterExt.activate()
8387
]);
@@ -87,22 +91,11 @@ suite('Smoke Tests', function () {
8791
}
8892
await jupyterExt.exports.openNotebook(notebook.uri, pythonEnv);
8993

90-
await vscode.commands.executeCommand<void>('notebook.execute');
91-
92-
await new Promise<void>((resolve) => {
93-
const disposable = vscode.workspace.onDidChangeNotebookDocument((e) => {
94-
e.cellChanges.forEach((change) => {
95-
if (
96-
change.outputs?.some((o) =>
97-
o.items.some((i) => Buffer.from(i.data).toString('utf-8').includes('Hello World'))
98-
)
99-
) {
100-
disposable.dispose();
101-
resolve();
102-
}
103-
});
104-
});
105-
});
94+
await Promise.all([
95+
vscode.commands.executeCommand<void>('notebook.execute'),
96+
waitForTextOutputV2(notebook.cellAt(0), 'Hello World', 0, false, disposableStore),
97+
waitForExecutionCompletedSuccessfullyV2(notebook.cellAt(0), disposableStore)
98+
]);
10699
}).timeout(timeoutForCellToRun);
107100

108101
test('Interactive window should always pick up current active interpreter', async function () {

0 commit comments

Comments
 (0)