From 7713b976b8a8941cf408eeeca201390dcc8420ec Mon Sep 17 00:00:00 2001
From: krassowski <5832902+krassowski@users.noreply.github.com>
Date: Sat, 24 Feb 2024 18:45:58 +0000
Subject: [PATCH] Fix diagnostics tests

---
 .../src/features/diagnostics/diagnostics.spec.ts           | 2 ++
 packages/jupyterlab-lsp/src/testutils.ts                   | 7 +++++--
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/packages/jupyterlab-lsp/src/features/diagnostics/diagnostics.spec.ts b/packages/jupyterlab-lsp/src/features/diagnostics/diagnostics.spec.ts
index fe0f0e208..4bcbe56cf 100644
--- a/packages/jupyterlab-lsp/src/features/diagnostics/diagnostics.spec.ts
+++ b/packages/jupyterlab-lsp/src/features/diagnostics/diagnostics.spec.ts
@@ -70,6 +70,7 @@ describe('Diagnostics', () => {
         editorExtensionRegistry: env.editorExtensionRegistry,
         themeManager: null
       });
+      env.featureManager.register(feature);
       await env.init();
     });
     afterEach(() => {
@@ -222,6 +223,7 @@ describe('Diagnostics', () => {
         editorExtensionRegistry: env.editorExtensionRegistry,
         themeManager: null
       });
+      env.featureManager.register(feature);
       await env.init();
     });
     afterEach(() => {
diff --git a/packages/jupyterlab-lsp/src/testutils.ts b/packages/jupyterlab-lsp/src/testutils.ts
index 8979ed360..7dd7c8797 100644
--- a/packages/jupyterlab-lsp/src/testutils.ts
+++ b/packages/jupyterlab-lsp/src/testutils.ts
@@ -24,7 +24,8 @@ import {
   DocumentConnectionManager,
   FeatureManager,
   ISocketConnectionOptions,
-  ILSPOptions
+  ILSPOptions,
+  ILSPFeatureManager
 } from '@jupyterlab/lsp';
 import { LSPConnection } from '@jupyterlab/lsp/lib/connection';
 import * as nbformat from '@jupyterlab/nbformat';
@@ -200,6 +201,7 @@ export abstract class TestEnvironment implements ITestEnvironment {
   documentOptions: VirtualDocument.IOptions;
   editorExtensionRegistry: EditorExtensionRegistry;
   editorServices: IEditorServices;
+  featureManager: ILSPFeatureManager;
 
   constructor(protected options?: TestEnvironment.IOptions) {
     this.editorExtensionRegistry = new EditorExtensionRegistry();
@@ -230,6 +232,7 @@ export abstract class TestEnvironment implements ITestEnvironment {
       ...this.getDefaults(),
       ...(options?.document || {})
     };
+    this.featureManager = new FeatureManager();
   }
 
   protected abstract createWidget(): IDocumentWidget;
@@ -280,7 +283,7 @@ export abstract class TestEnvironment implements ITestEnvironment {
       docRegistry,
       connectionManager: this.connectionManager,
       codeOverridesManager: overridesManager,
-      featureManager: new FeatureManager(),
+      featureManager: this.featureManager,
       foreignCodeExtractorsManager: foreignCodeExtractors,
       translator: nullTranslator
     });