/*--------------------------------------------------------------------------------------------- * 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 * as sinon from 'sinon'; import { CancellationTokenSource, Disposable, EventEmitter, ExtensionContext, NotebookCellKind, NotebookDocumentChangeEvent, NotebookDocumentWillSaveEvent, NotebookEdit, NotebookRange, TextDocument, TextDocumentSaveReason, workspace, type CancellationToken, type NotebookCell, type NotebookDocument, type WorkspaceEdit, type WorkspaceEditMetadata } from 'vscode'; import { activate } from '../notebookModelStoreSync'; suite(`Notebook Model Store Sync`, () => { let disposables: Disposable[] = []; let onDidChangeNotebookDocument: EventEmitter; let onWillSaveNotebookDocument: AsyncEmitter; let notebook: NotebookDocument; let token: CancellationTokenSource; let editsApplied: WorkspaceEdit[] = []; let pendingPromises: Promise[] = []; let cellMetadataUpdates: NotebookEdit[] = []; let applyEditStub: sinon.SinonStub<[edit: WorkspaceEdit, metadata?: WorkspaceEditMetadata | undefined], Thenable>; setup(() => { disposables = []; notebook = { notebookType: '', metadata: {} } as NotebookDocument; token = new CancellationTokenSource(); disposables.push(token); sinon.stub(notebook, 'notebookType').get(() => 'jupyter-notebook'); applyEditStub = sinon.stub(workspace, 'applyEdit').callsFake((edit: WorkspaceEdit) => { editsApplied.push(edit); return Promise.resolve(true); }); const context = { subscriptions: [] as Disposable[] } as ExtensionContext; onDidChangeNotebookDocument = new EventEmitter(); disposables.push(onDidChangeNotebookDocument); onWillSaveNotebookDocument = new AsyncEmitter(); const stub = sinon.stub(NotebookEdit, 'updateCellMetadata').callsFake((index, metadata) => { const edit = stub.wrappedMethod.call(NotebookEdit, index, metadata); cellMetadataUpdates.push(edit); return edit; } ); sinon.stub(workspace, 'onDidChangeNotebookDocument').callsFake(cb => onDidChangeNotebookDocument.event(cb) ); sinon.stub(workspace, 'onWillSaveNotebookDocument').callsFake(cb => onWillSaveNotebookDocument.event(cb) ); activate(context); }); teardown(async () => { await Promise.allSettled(pendingPromises); editsApplied = []; pendingPromises = []; cellMetadataUpdates = []; disposables.forEach(d => d.dispose()); disposables = []; sinon.restore(); }); test('Empty cell will not result in any updates', async () => { const e: NotebookDocumentChangeEvent = { notebook, metadata: undefined, contentChanges: [], cellChanges: [] }; onDidChangeNotebookDocument.fire(e); assert.strictEqual(editsApplied.length, 0); }); test('Adding cell for non Jupyter Notebook will not result in any updates', async () => { sinon.stub(notebook, 'notebookType').get(() => 'some-other-type'); const cell: NotebookCell = { document: {} as unknown as TextDocument, executionSummary: {}, index: 0, kind: NotebookCellKind.Code, metadata: {}, notebook, outputs: [] }; const e: NotebookDocumentChangeEvent = { notebook, metadata: undefined, contentChanges: [ { range: new NotebookRange(0, 0), removedCells: [], addedCells: [cell] } ], cellChanges: [] }; onDidChangeNotebookDocument.fire(e); assert.strictEqual(editsApplied.length, 0); assert.strictEqual(cellMetadataUpdates.length, 0); }); test('Adding cell to nbformat 4.2 notebook will result in adding empty metadata', async () => { sinon.stub(notebook, 'metadata').get(() => ({ nbformat: 4, nbformat_minor: 2 })); const cell: NotebookCell = { document: {} as unknown as TextDocument, executionSummary: {}, index: 0, kind: NotebookCellKind.Code, metadata: {}, notebook, outputs: [] }; const e: NotebookDocumentChangeEvent = { notebook, metadata: undefined, contentChanges: [ { range: new NotebookRange(0, 0), removedCells: [], addedCells: [cell] } ], cellChanges: [] }; onDidChangeNotebookDocument.fire(e); assert.strictEqual(editsApplied.length, 1); assert.strictEqual(cellMetadataUpdates.length, 1); const newMetadata = cellMetadataUpdates[0].newCellMetadata; assert.deepStrictEqual(newMetadata, { execution_count: null, metadata: {} }); }); test('Added cell will have a cell id if nbformat is 4.5', async () => { sinon.stub(notebook, 'metadata').get(() => ({ nbformat: 4, nbformat_minor: 5 })); const cell: NotebookCell = { document: {} as unknown as TextDocument, executionSummary: {}, index: 0, kind: NotebookCellKind.Code, metadata: {}, notebook, outputs: [] }; const e: NotebookDocumentChangeEvent = { notebook, metadata: undefined, contentChanges: [ { range: new NotebookRange(0, 0), removedCells: [], addedCells: [cell] } ], cellChanges: [] }; onDidChangeNotebookDocument.fire(e); assert.strictEqual(editsApplied.length, 1); assert.strictEqual(cellMetadataUpdates.length, 1); const newMetadata = cellMetadataUpdates[0].newCellMetadata || {}; assert.strictEqual(Object.keys(newMetadata).length, 3); assert.deepStrictEqual(newMetadata.execution_count, null); assert.deepStrictEqual(newMetadata.metadata, {}); assert.ok(newMetadata.id); }); test('Do not add cell id if one already exists', async () => { sinon.stub(notebook, 'metadata').get(() => ({ nbformat: 4, nbformat_minor: 5 })); const cell: NotebookCell = { document: {} as unknown as TextDocument, executionSummary: {}, index: 0, kind: NotebookCellKind.Code, metadata: { id: '1234' }, notebook, outputs: [] }; const e: NotebookDocumentChangeEvent = { notebook, metadata: undefined, contentChanges: [ { range: new NotebookRange(0, 0), removedCells: [], addedCells: [cell] } ], cellChanges: [] }; onDidChangeNotebookDocument.fire(e); assert.strictEqual(editsApplied.length, 1); assert.strictEqual(cellMetadataUpdates.length, 1); const newMetadata = cellMetadataUpdates[0].newCellMetadata || {}; assert.strictEqual(Object.keys(newMetadata).length, 3); assert.deepStrictEqual(newMetadata.execution_count, null); assert.deepStrictEqual(newMetadata.metadata, {}); assert.strictEqual(newMetadata.id, '1234'); }); test('Do not perform any updates if cell id and metadata exists', async () => { sinon.stub(notebook, 'metadata').get(() => ({ nbformat: 4, nbformat_minor: 5 })); const cell: NotebookCell = { document: {} as unknown as TextDocument, executionSummary: {}, index: 0, kind: NotebookCellKind.Code, metadata: { id: '1234', metadata: {} }, notebook, outputs: [] }; const e: NotebookDocumentChangeEvent = { notebook, metadata: undefined, contentChanges: [ { range: new NotebookRange(0, 0), removedCells: [], addedCells: [cell] } ], cellChanges: [] }; onDidChangeNotebookDocument.fire(e); assert.strictEqual(editsApplied.length, 0); assert.strictEqual(cellMetadataUpdates.length, 0); }); test('Store language id in custom metadata, whilst preserving existing metadata', async () => { sinon.stub(notebook, 'metadata').get(() => ({ nbformat: 4, nbformat_minor: 5, metadata: { language_info: { name: 'python' } } })); const cell: NotebookCell = { document: { languageId: 'javascript' } as unknown as TextDocument, executionSummary: {}, index: 0, kind: NotebookCellKind.Code, metadata: { id: '1234', metadata: { collapsed: true, scrolled: true } }, notebook, outputs: [] }; const e: NotebookDocumentChangeEvent = { notebook, metadata: undefined, contentChanges: [], cellChanges: [ { cell, document: { languageId: 'javascript' } as unknown as TextDocument, metadata: undefined, outputs: undefined, executionSummary: undefined } ] }; onDidChangeNotebookDocument.fire(e); assert.strictEqual(editsApplied.length, 1); assert.strictEqual(cellMetadataUpdates.length, 1); const newMetadata = cellMetadataUpdates[0].newCellMetadata || {}; assert.strictEqual(Object.keys(newMetadata).length, 3); assert.deepStrictEqual(newMetadata.execution_count, null); assert.deepStrictEqual(newMetadata.metadata, { collapsed: true, scrolled: true, vscode: { languageId: 'javascript' } }); assert.strictEqual(newMetadata.id, '1234'); }); test('No changes when language is javascript', async () => { sinon.stub(notebook, 'metadata').get(() => ({ nbformat: 4, nbformat_minor: 5, metadata: { language_info: { name: 'javascript' } } })); const cell: NotebookCell = { document: { languageId: 'javascript' } as unknown as TextDocument, executionSummary: {}, index: 0, kind: NotebookCellKind.Code, metadata: { id: '1234', metadata: { collapsed: true, scrolled: true } }, notebook, outputs: [] }; const e: NotebookDocumentChangeEvent = { notebook, metadata: undefined, contentChanges: [], cellChanges: [ { cell, document: undefined, metadata: undefined, outputs: undefined, executionSummary: undefined } ] }; onDidChangeNotebookDocument.fire(e); assert.strictEqual(editsApplied.length, 0); assert.strictEqual(cellMetadataUpdates.length, 0); }); test('Remove language from metadata when cell language matches kernel language', async () => { sinon.stub(notebook, 'metadata').get(() => ({ nbformat: 4, nbformat_minor: 5, metadata: { language_info: { name: 'javascript' } } })); const cell: NotebookCell = { document: { languageId: 'javascript' } as unknown as TextDocument, executionSummary: {}, index: 0, kind: NotebookCellKind.Code, metadata: { id: '1234', metadata: { vscode: { languageId: 'python' }, collapsed: true, scrolled: true } }, notebook, outputs: [] }; const e: NotebookDocumentChangeEvent = { notebook, metadata: undefined, contentChanges: [], cellChanges: [ { cell, document: { languageId: 'javascript' } as unknown as TextDocument, metadata: undefined, outputs: undefined, executionSummary: undefined } ] }; onDidChangeNotebookDocument.fire(e); assert.strictEqual(editsApplied.length, 1); assert.strictEqual(cellMetadataUpdates.length, 1); const newMetadata = cellMetadataUpdates[0].newCellMetadata || {}; assert.strictEqual(Object.keys(newMetadata).length, 3); assert.deepStrictEqual(newMetadata.execution_count, null); assert.deepStrictEqual(newMetadata.metadata, { collapsed: true, scrolled: true }); assert.strictEqual(newMetadata.id, '1234'); }); test('Update language in metadata', async () => { sinon.stub(notebook, 'metadata').get(() => ({ nbformat: 4, nbformat_minor: 5, metadata: { language_info: { name: 'javascript' } } })); const cell: NotebookCell = { document: { languageId: 'powershell' } as unknown as TextDocument, executionSummary: {}, index: 0, kind: NotebookCellKind.Code, metadata: { id: '1234', metadata: { vscode: { languageId: 'python' }, collapsed: true, scrolled: true } }, notebook, outputs: [] }; const e: NotebookDocumentChangeEvent = { notebook, metadata: undefined, contentChanges: [], cellChanges: [ { cell, document: { languageId: 'powershell' } as unknown as TextDocument, metadata: undefined, outputs: undefined, executionSummary: undefined } ] }; onDidChangeNotebookDocument.fire(e); assert.strictEqual(editsApplied.length, 1); assert.strictEqual(cellMetadataUpdates.length, 1); const newMetadata = cellMetadataUpdates[0].newCellMetadata || {}; assert.strictEqual(Object.keys(newMetadata).length, 3); assert.deepStrictEqual(newMetadata.execution_count, null); assert.deepStrictEqual(newMetadata.metadata, { collapsed: true, scrolled: true, vscode: { languageId: 'powershell' } }); assert.strictEqual(newMetadata.id, '1234'); }); test('Will save event without any changes', async () => { await onWillSaveNotebookDocument.fireAsync({ notebook, reason: TextDocumentSaveReason.Manual }, token.token); }); test('Wait for pending updates to complete when saving', async () => { let resolveApplyEditPromise: (value: boolean) => void; const promise = new Promise((resolve) => resolveApplyEditPromise = resolve); applyEditStub.restore(); sinon.stub(workspace, 'applyEdit').callsFake((edit: WorkspaceEdit) => { editsApplied.push(edit); return promise; }); const cell: NotebookCell = { document: {} as unknown as TextDocument, executionSummary: {}, index: 0, kind: NotebookCellKind.Code, metadata: {}, notebook, outputs: [] }; const e: NotebookDocumentChangeEvent = { notebook, metadata: undefined, contentChanges: [ { range: new NotebookRange(0, 0), removedCells: [], addedCells: [cell] } ], cellChanges: [] }; onDidChangeNotebookDocument.fire(e); assert.strictEqual(editsApplied.length, 1); assert.strictEqual(cellMetadataUpdates.length, 1); // Try to save. let saveCompleted = false; const saved = onWillSaveNotebookDocument.fireAsync({ notebook, reason: TextDocumentSaveReason.Manual }, token.token); saved.finally(() => saveCompleted = true); await new Promise((resolve) => setTimeout(resolve, 10)); // Verify we have not yet completed saving. assert.strictEqual(saveCompleted, false); resolveApplyEditPromise!(true); await new Promise((resolve) => setTimeout(resolve, 1)); // Should have completed saving. saved.finally(() => saveCompleted = true); }); interface IWaitUntil { token: CancellationToken; waitUntil(thenable: Promise): void; } interface IWaitUntil { token: CancellationToken; waitUntil(thenable: Promise): void; } type IWaitUntilData = Omit, 'token'>; class AsyncEmitter { private listeners: ((d: T) => void)[] = []; get event(): (listener: (e: T) => any, thisArgs?: any, disposables?: Disposable[]) => Disposable { return (listener, thisArgs, _disposables) => { this.listeners.push(listener.bind(thisArgs)); return { dispose: () => { // } }; }; } dispose() { this.listeners = []; } async fireAsync(data: IWaitUntilData, token: CancellationToken): Promise { if (!this.listeners.length) { return; } const promises: Promise[] = []; this.listeners.forEach(cb => { const event = { ...data, token, waitUntil: (thenable: Promise) => { promises.push(thenable); } } as T; cb(event); }); await Promise.all(promises); } } });