@@ -43,7 +43,8 @@ import {
43
43
env ,
44
44
languages ,
45
45
window ,
46
- workspace
46
+ workspace ,
47
+ type CancellationToken
47
48
} from 'vscode' ;
48
49
import { DebugProtocol } from 'vscode-debugprotocol' ;
49
50
import {
@@ -78,11 +79,11 @@ import {
78
79
defaultNotebookFormat ,
79
80
isWebExtension
80
81
} from '../../../platform/common/constants' ;
81
- import { dispose } from '../../../platform/common/utils/lifecycle' ;
82
+ import { dispose , type DisposableStore } from '../../../platform/common/utils/lifecycle' ;
82
83
import { getDisplayPath } from '../../../platform/common/platform/fs-paths' ;
83
84
import { IFileSystem , IPlatformService } from '../../../platform/common/platform/types' ;
84
85
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' ;
86
87
import { DataScience } from '../../../platform/common/utils/localize' ;
87
88
import { isWeb } from '../../../platform/common/utils/misc' ;
88
89
import { openAndShowNotebook } from '../../../platform/common/utils/notebooks' ;
@@ -104,6 +105,7 @@ import {
104
105
NotebookCellExecutionState ,
105
106
notebookCellExecutions
106
107
} from '../../../platform/notebooks/cellExecutionStateService' ;
108
+ import { raceCancellationError } from '../../../platform/common/cancellation' ;
107
109
108
110
// Running in Conda environments, things can be a little slower.
109
111
export const defaultNotebookTestTimeout = 60_000 ;
@@ -982,6 +984,40 @@ export async function waitForExecutionCompletedSuccessfully(
982
984
await sleep ( 100 ) ;
983
985
}
984
986
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
+
985
1021
export async function waitForCompletions (
986
1022
completionProvider : CompletionItemProvider ,
987
1023
cell : NotebookCell ,
@@ -1220,6 +1256,38 @@ export async function waitForTextOutput(
1220
1256
. join ( ',\n' ) } `
1221
1257
) ;
1222
1258
}
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
+ }
1223
1291
export function assertNotHasTextOutputInVSCode ( cell : NotebookCell , text : string , index : number , isExactMatch = true ) {
1224
1292
const cellOutputs = cell . outputs ;
1225
1293
assert . ok ( cellOutputs , 'No output' ) ;
0 commit comments