Wait for schemas to be loaded

Author: satyakigh

src/datastore/LMDB.ts

@@ -12,7 +12,6 @@ import { stats } from './lmdb/Stats';
 import { encryptionStrategy } from './lmdb/Utils';
 
 export class LMDBStoreFactory implements DataStoreFactory {
-    private isClosed: boolean = false;
     private readonly log = LoggerFactory.getLogger('LMDB.Global');
     @Telemetry({ scope: 'LMDB.Global' }) private readonly telemetry!: ScopedTelemetry;
 
@@ -87,18 +86,14 @@ export class LMDBStoreFactory implements DataStoreFactory {
     }
 
     async close(): Promise<void> {
-        if (this.isClosed) {
-            return;
-        }
+        clearInterval(this.metricsInterval);
+        clearTimeout(this.timeout);
 
         // Clear the stores map but don't close individual stores
         // LMDB will close them when we close the environment
-        clearInterval(this.metricsInterval);
-        clearTimeout(this.timeout);
         this.stores.clear();
         await this.env.flushed;
         await this.env.close();
-        this.isClosed = true;
     }
 
     private cleanupOldVersions(): void {

src/server/CfnServer.ts

@@ -61,7 +61,6 @@ import { ServerComponents } from './ServerComponents';
 
 const log = LoggerFactory.getLogger('CfnServer');
 export class CfnServer {
-    private isClosed: boolean = false;
     private readonly components: ServerComponents;
 
     constructor(
@@ -290,10 +289,6 @@ export class CfnServer {
     }
 
     async close(): Promise<void> {
-        if (this.isClosed) {
-            return;
-        }
         await closeSafely(this.providers, this.external, this.core);
-        this.isClosed = true;
     }
 }

tst/unit/utils/wheel-download.test.ts tst/unit/tools/wheel-download.test.tstst/unit/utils/wheel-download.test.ts renamed to tst/unit/tools/wheel-download.test.ts

@@ -53,6 +53,7 @@ import { MultiDataStoreFactoryProvider } from '../../src/datastore/DataStore';
 import { FeatureFlagProvider } from '../../src/featureFlag/FeatureFlagProvider';
 import { LspCapabilities } from '../../src/protocol/LspCapabilities';
 import { LspConnection } from '../../src/protocol/LspConnection';
+import { SamStoreKey } from '../../src/schema/SamSchemas';
 import { SchemaRetriever } from '../../src/schema/SchemaRetriever';
 import { SchemaStore } from '../../src/schema/SchemaStore';
 import { CfnExternal } from '../../src/server/CfnExternal';
@@ -61,39 +62,23 @@ import { CfnLspProviders } from '../../src/server/CfnLspProviders';
 import { CfnServer } from '../../src/server/CfnServer';
 import { AwsMetadata, ExtendedInitializeParams } from '../../src/server/InitParams';
 import { RelationshipSchemaService } from '../../src/services/RelationshipSchemaService';
+import { DefaultSettings } from '../../src/settings/Settings';
 import { LoggerFactory } from '../../src/telemetry/LoggerFactory';
 import { Closeable } from '../../src/utils/Closeable';
 import { ExtensionName } from '../../src/utils/ExtensionConfig';
 import { createMockCfnLintService } from './MockServerComponents';
 import { getTestPrivateSchemas, samFileType, SamSchemaFiles, schemaFileType, Schemas } from './SchemaUtils';
 import { flushAllPromises, WaitFor } from './Utils';
 
+type TestExtensionConfig = {
+    id?: string;
+    initializeParams?: Partial<ExtendedInitializeParams>;
+    workspaceConfig?: Record<string, unknown>[];
+};
+
 export class TestExtension implements Closeable {
-    private readonly id = v4();
-    private readonly rootDir = join(process.cwd(), 'node_modules', '.cache', 'e2e-tests', this.id);
-    private readonly awsMetadata: AwsMetadata = {
-        clientInfo: {
-            extension: {
-                name: `Test ${ExtensionName}`,
-                version: '1.0.0-test',
-            },
-            clientId: this.id,
-        },
-        encryption: {
-            key: randomBytes(32).toString('base64'),
-            mode: 'JWT',
-        },
-    };
-    private readonly initializeParams: ExtendedInitializeParams = {
-        processId: process.pid,
-        rootUri: null,
-        capabilities: {},
-        clientInfo: this.awsMetadata.clientInfo?.extension,
-        workspaceFolders: [],
-        initializationOptions: {
-            aws: this.awsMetadata,
-        },
-    };
+    private readonly awsMetadata: AwsMetadata;
+    private readonly initializeParams: ExtendedInitializeParams;
 
     private readonly readStream = new PassThrough();
     private readonly writeStream = new PassThrough();
@@ -108,18 +93,45 @@ export class TestExtension implements Closeable {
     providers!: CfnLspProviders;
     server!: CfnServer;
 
-    private lspClientReady = false;
-    private lspServerReady = false;
+    private isReady = false;
+
+    constructor(config: TestExtensionConfig = {}) {
+        const id = config.id ?? v4();
+        const rootDir = join(process.cwd(), 'node_modules', '.cache', 'e2e-tests', id);
+
+        this.awsMetadata = {
+            clientInfo: {
+                extension: {
+                    name: `Test ${ExtensionName}`,
+                    version: '1.0.0-test',
+                },
+                clientId: id,
+            },
+            encryption: {
+                key: randomBytes(32).toString('base64'),
+                mode: 'JWT',
+            },
+        };
+        this.initializeParams = {
+            processId: process.pid,
+            rootUri: null,
+            capabilities: {},
+            clientInfo: this.awsMetadata.clientInfo?.extension,
+            workspaceFolders: [],
+            initializationOptions: {
+                aws: this.awsMetadata,
+            },
+            ...config.initializeParams,
+        };
 
-    constructor() {
         this.serverConnection = new LspConnection(
             createConnection(new StreamMessageReader(this.readStream), new StreamMessageWriter(this.writeStream)),
             {
                 onInitialize: (params) => {
                     const lsp = this.serverConnection.components;
                     LoggerFactory.reconfigure('warn');
 
-                    const dataStoreFactory = new MultiDataStoreFactoryProvider(this.rootDir);
+                    const dataStoreFactory = new MultiDataStoreFactoryProvider(rootDir);
                     this.core = new CfnInfraCore(lsp, params, {
                         dataStoreFactory,
                     });
@@ -154,22 +166,15 @@ export class TestExtension implements Closeable {
                     this.server = new CfnServer(lsp, this.core, this.external, this.providers);
                     return LspCapabilities;
                 },
-                onInitialized: (params) => {
-                    this.server.initialized(params);
-                    this.lspServerReady = true;
-                },
-                onShutdown: () => {
-                    return this.server.close();
-                },
-                onExit: () => {
-                    return this.server.close();
-                },
+                onInitialized: (params) => this.server.initialized(params),
+                onShutdown: () => this.server.close(),
+                onExit: () => this.server.close(),
             },
         );
 
         // Handle workspace/configuration requests from the server
         this.clientConnection.onRequest('workspace/configuration', () => {
-            return [{}]; // Return empty configuration
+            return config.workspaceConfig ?? [{}];
         });
 
         this.serverConnection.listen();
@@ -185,40 +190,39 @@ export class TestExtension implements Closeable {
     }
 
     async ready() {
-        if (!this.lspClientReady) {
+        if (!this.isReady) {
             await this.clientConnection.sendRequest(InitializeRequest.type, this.initializeParams);
             await this.clientConnection.sendNotification(InitializedNotification.type, {});
-            this.lspClientReady = true;
-        }
 
-        await WaitFor.waitFor(() => {
-            if (!this.lspServerReady) {
-                throw new Error('Server is not ready yet');
-            }
-        }, 5000);
+            await WaitFor.waitFor(() => {
+                const store = this.external.schemaStore;
+                const pbSchemas = store?.publicSchemas?.get(DefaultSettings.profile.region);
+                const samSchemas = store?.samSchemas?.get(SamStoreKey);
 
-        await flushAllPromises();
+                if (pbSchemas === undefined || samSchemas === undefined) {
+                    throw new Error('Schemas not loaded yet');
+                }
+            }, 1_000);
+
+            await flushAllPromises();
+            this.isReady = true;
+        }
     }
 
     async send(method: string, params: any) {
         await this.ready();
-        const value = await this.clientConnection.sendRequest(method, params);
-        await wait(100);
-        return value;
+        return await this.clientConnection.sendRequest(method, params);
     }
 
     async notify(method: string, params: any) {
         await this.ready();
-        const value = await this.clientConnection.sendNotification(method, params);
-        await wait(100);
-        return value;
+        return await this.clientConnection.sendNotification(method, params);
     }
 
     async close() {
         await this.clientConnection.sendRequest(ShutdownRequest.type);
         await this.clientConnection.sendNotification(ExitNotification.type);
         this.clientConnection.dispose();
-        await flushAllPromises();
     }
 
     // ====================================================================
@@ -227,18 +231,22 @@ export class TestExtension implements Closeable {
 
     async openDocument(params: DidOpenTextDocumentParams) {
         await this.notify(DidOpenTextDocumentNotification.method, params);
+        await wait(10);
     }
 
     async changeDocument(params: DidChangeTextDocumentParams) {
         await this.notify(DidChangeTextDocumentNotification.method, params);
+        await wait(10);
     }
 
     async closeDocument(params: DidCloseTextDocumentParams) {
         await this.notify(DidCloseTextDocumentNotification.method, params);
+        await wait(10);
     }
 
-    saveDocument(params: DidSaveTextDocumentParams) {
-        return this.notify(DidSaveTextDocumentNotification.method, params);
+    async saveDocument(params: DidSaveTextDocumentParams) {
+        await this.notify(DidSaveTextDocumentNotification.method, params);
+        await wait(10);
     }
 
     completion(params: CompletionParams) {

tst/utils/TestExtension.ts

@@ -1,13 +1,13 @@
-export function flushAllPromises() {
-    return new Promise((resolve) => {
-        setTimeout(resolve, 0);
-    });
+import { setImmediate } from 'node:timers/promises';
+
+export async function flushAllPromises() {
+    await setImmediate();
 }
 
 export class WaitFor {
     constructor(
-        private readonly maxWaitMs: number = 25,
-        private readonly delayIntervalMs: number = 1,
+        private readonly maxWaitMs: number = 100,
+        private readonly delayIntervalMs: number = 10,
         private readonly throwableTypesToExpect: (new (...args: any[]) => Error)[] = [Error],
     ) {}
 
@@ -33,7 +33,11 @@ export class WaitFor {
         throw lastError!;
     }
 
-    static async waitFor(code: () => void | Promise<void>, timeoutMs: number = 100): Promise<void> {
-        await new WaitFor(timeoutMs).wait(code);
+    static async waitFor(
+        code: () => void | Promise<void>,
+        timeoutMs: number = 100,
+        intervalMs: number = 10,
+    ): Promise<void> {
+        await new WaitFor(timeoutMs, intervalMs).wait(code);
     }
 }