diff --git a/src/datastore/DataStore.ts b/src/datastore/DataStore.ts index 995510e3..deff8206 100644 --- a/src/datastore/DataStore.ts +++ b/src/datastore/DataStore.ts @@ -14,7 +14,6 @@ export enum StoreName { public_schemas = 'public_schemas', sam_schemas = 'sam_schemas', private_schemas = 'private_schemas', - combined_schemas = 'combined_schemas', } export interface DataStore { diff --git a/src/handlers/Initialize.ts b/src/handlers/Initialize.ts index 42012dac..03d82f6a 100644 --- a/src/handlers/Initialize.ts +++ b/src/handlers/Initialize.ts @@ -6,17 +6,16 @@ const logger = LoggerFactory.getLogger('InitializedHandler'); export function initializedHandler(workspace: LspWorkspace, components: ServerComponents): () => void { return (): void => { - // Sync configuration from LSP workspace first, then initialize CfnLintService components.settingsManager .syncConfiguration() .then(() => { + components.schemaRetriever.initialize(); return components.cfnLintService.initialize(); }) .then(async () => { // Process folders sequentially to avoid overwhelming the system for (const folder of workspace.getAllWorkspaceFolders()) { try { - // Properly await the async mountFolder method await components.cfnLintService.mountFolder(folder); } catch (error) { logger.error(error, `Failed to mount folder ${folder.name}`); diff --git a/src/schema/CombinedSchemas.ts b/src/schema/CombinedSchemas.ts index cf173df3..6004650a 100644 --- a/src/schema/CombinedSchemas.ts +++ b/src/schema/CombinedSchemas.ts @@ -1,15 +1,17 @@ -import { LoggerFactory } from '../telemetry/LoggerFactory'; import { PrivateSchemas, PrivateSchemasType } from './PrivateSchemas'; import { RegionalSchemas, RegionalSchemasType } from './RegionalSchemas'; import { ResourceSchema } from './ResourceSchema'; import { SamSchemas, SamSchemasType } from './SamSchemas'; export class CombinedSchemas { - private static readonly log = LoggerFactory.getLogger('CombinedSchemas'); readonly numSchemas: number; readonly schemas: Map; - constructor(regionalSchemas?: RegionalSchemas, privateSchemas?: PrivateSchemas, samSchemas?: SamSchemas) { + constructor( + readonly regionalSchemas?: RegionalSchemas, + readonly privateSchemas?: PrivateSchemas, + readonly samSchemas?: SamSchemas, + ) { this.schemas = new Map([ ...(privateSchemas?.schemas ?? []), ...(regionalSchemas?.schemas ?? []), @@ -22,14 +24,10 @@ export class CombinedSchemas { regionalSchemas?: RegionalSchemasType, privateSchemas?: PrivateSchemasType, samSchemas?: SamSchemasType, - ) { + ): CombinedSchemas { const regionalSchema = regionalSchemas === undefined ? undefined : RegionalSchemas.from(regionalSchemas); const privateSchema = privateSchemas === undefined ? undefined : PrivateSchemas.from(privateSchemas); const samSchema = samSchemas === undefined ? undefined : SamSchemas.from(samSchemas); - - CombinedSchemas.log.info( - `Combined schemas from public=${regionalSchemas?.schemas.length}, private=${privateSchema?.schemas.size}, SAM=${samSchema?.schemas.size}`, - ); return new CombinedSchemas(regionalSchema, privateSchema, samSchema); } } diff --git a/src/schema/GetSamSchemaTask.ts b/src/schema/GetSamSchemaTask.ts index eb3e5d5a..bc227a89 100644 --- a/src/schema/GetSamSchemaTask.ts +++ b/src/schema/GetSamSchemaTask.ts @@ -1,5 +1,5 @@ -import { Logger } from 'pino'; import { DataStore } from '../datastore/DataStore'; +import { LoggerFactory } from '../telemetry/LoggerFactory'; import { Measure } from '../telemetry/TelemetryDecorator'; import { downloadJson } from '../utils/RemoteDownload'; import { GetSchemaTask } from './GetSchemaTask'; @@ -7,12 +7,14 @@ import { SamSchemas, SamSchemasType, SamStoreKey } from './SamSchemas'; import { CloudFormationResourceSchema, SamSchema, SamSchemaTransformer } from './SamSchemaTransformer'; export class GetSamSchemaTask extends GetSchemaTask { + private readonly logger = LoggerFactory.getLogger(GetSamSchemaTask); + constructor(private readonly getSamSchemas: () => Promise>) { super(); } @Measure({ name: 'getSchemas' }) - protected override async runImpl(dataStore: DataStore, logger?: Logger): Promise { + protected override async runImpl(dataStore: DataStore): Promise { try { const resourceSchemas = await this.getSamSchemas(); @@ -32,9 +34,9 @@ export class GetSamSchemaTask extends GetSchemaTask { await dataStore.put(SamStoreKey, samSchemasData); - logger?.info(`${resourceSchemas.size} SAM schemas downloaded and stored`); + this.logger.info(`${resourceSchemas.size} SAM schemas downloaded and stored`); } catch (error) { - logger?.error(error, 'Failed to download SAM schema'); + this.logger.error(error, 'Failed to download SAM schemas'); throw error; } } diff --git a/src/schema/GetSchemaTask.ts b/src/schema/GetSchemaTask.ts index 1a57e7c6..5e9a800c 100644 --- a/src/schema/GetSchemaTask.ts +++ b/src/schema/GetSchemaTask.ts @@ -1,25 +1,27 @@ import { DescribeTypeOutput } from '@aws-sdk/client-cloudformation'; -import { Logger } from 'pino'; import { AwsCredentials } from '../auth/AwsCredentials'; import { DataStore } from '../datastore/DataStore'; import { CfnService } from '../services/CfnService'; +import { LoggerFactory } from '../telemetry/LoggerFactory'; import { ScopedTelemetry } from '../telemetry/ScopedTelemetry'; import { Measure, Telemetry } from '../telemetry/TelemetryDecorator'; import { AwsRegion } from '../utils/Region'; import { downloadFile } from '../utils/RemoteDownload'; -import { PrivateSchemas, PrivateSchemasType } from './PrivateSchemas'; +import { PrivateSchemas, PrivateSchemasType, PrivateStoreKey } from './PrivateSchemas'; import { RegionalSchemas, RegionalSchemasType, SchemaFileType } from './RegionalSchemas'; import { cfnResourceSchemaLink, unZipFile } from './RemoteSchemaHelper'; export abstract class GetSchemaTask { - protected abstract runImpl(dataStore: DataStore, logger?: Logger): Promise; + protected abstract runImpl(dataStore: DataStore): Promise; - async run(dataStore: DataStore, logger?: Logger) { - await this.runImpl(dataStore, logger); + async run(dataStore: DataStore) { + await this.runImpl(dataStore); } } export class GetPublicSchemaTask extends GetSchemaTask { + private readonly logger = LoggerFactory.getLogger(GetPublicSchemaTask); + @Telemetry() private readonly telemetry!: ScopedTelemetry; @@ -35,9 +37,20 @@ export class GetPublicSchemaTask extends GetSchemaTask { } @Measure({ name: 'getSchemas' }) - protected override async runImpl(dataStore: DataStore, logger?: Logger) { + protected override async runImpl(dataStore: DataStore) { + this.telemetry.count(`getSchemas.maxAttempt.fault`, 0, { + attributes: { + region: this.region, + }, + }); + if (this.attempts >= GetPublicSchemaTask.MaxAttempts) { - logger?.error(`Reached max attempts for retrieving schemas for ${this.region} without success`); + this.telemetry.count(`getSchemas.maxAttempt.fault`, 1, { + attributes: { + region: this.region, + }, + }); + this.logger.error(`Reached max attempts for retrieving schemas for ${this.region} without success`); return; } @@ -53,44 +66,35 @@ export class GetPublicSchemaTask extends GetSchemaTask { }; await dataStore.put(this.region, value); - logger?.info(`${schemas.length} public schemas retrieved for ${this.region}`); + this.logger.info(`${schemas.length} public schemas retrieved for ${this.region}`); } } export class GetPrivateSchemasTask extends GetSchemaTask { - private readonly processedProfiles = new Set(); + private readonly logger = LoggerFactory.getLogger(GetPrivateSchemasTask); - constructor( - private readonly getSchemas: () => Promise, - private readonly getProfile: () => string, - ) { + constructor(private readonly getSchemas: () => Promise) { super(); } @Measure({ name: 'getSchemas' }) - protected override async runImpl(dataStore: DataStore, logger?: Logger) { + protected override async runImpl(dataStore: DataStore) { try { - const profile = this.getProfile(); - if (this.processedProfiles.has(profile)) { - return; - } - const schemas: DescribeTypeOutput[] = await this.getSchemas(); const value: PrivateSchemasType = { version: PrivateSchemas.V1, - identifier: profile, + identifier: PrivateStoreKey, schemas: schemas, firstCreatedMs: Date.now(), lastModifiedMs: Date.now(), }; - await dataStore.put(profile, value); + await dataStore.put(PrivateStoreKey, value); - this.processedProfiles.add(profile); - logger?.info(`${schemas.length} private schemas retrieved`); + this.logger.info(`${schemas.length} private schemas retrieved`); } catch (error) { - logger?.error(error, `Failed to get private schemas`); + this.logger.error(error, 'Failed to get private schemas'); throw error; } } diff --git a/src/schema/GetSchemaTaskManager.ts b/src/schema/GetSchemaTaskManager.ts index 4fc2c4a5..29c02701 100644 --- a/src/schema/GetSchemaTaskManager.ts +++ b/src/schema/GetSchemaTaskManager.ts @@ -1,132 +1,73 @@ import { DescribeTypeOutput } from '@aws-sdk/client-cloudformation'; -import { SettingsConfigurable, ISettingsSubscriber, SettingsSubscription } from '../settings/ISettingsSubscriber'; -import { DefaultSettings } from '../settings/Settings'; import { LoggerFactory } from '../telemetry/LoggerFactory'; -import { Closeable } from '../utils/Closeable'; -import { AwsRegion } from '../utils/Region'; +import { AwsRegion, getRegion } from '../utils/Region'; import { GetSamSchemaTask } from './GetSamSchemaTask'; import { GetPrivateSchemasTask, GetPublicSchemaTask } from './GetSchemaTask'; import { SchemaFileType } from './RegionalSchemas'; import { CloudFormationResourceSchema } from './SamSchemaTransformer'; import { SchemaStore } from './SchemaStore'; -const TenSeconds = 10 * 1000; -const OneHour = 60 * 60 * 1000; - -export class GetSchemaTaskManager implements SettingsConfigurable, Closeable { +export class GetSchemaTaskManager { + private readonly processedRegions = new Set(); private readonly tasks: GetPublicSchemaTask[] = []; private readonly privateTask: GetPrivateSchemasTask; private readonly samTask: GetSamSchemaTask; - private settingsSubscription?: SettingsSubscription; private readonly log = LoggerFactory.getLogger(GetSchemaTaskManager); - - private isRunning: boolean = false; - - private readonly timeout: NodeJS.Timeout; - private readonly interval: NodeJS.Timeout; + private isRunning = false; constructor( private readonly schemas: SchemaStore, private readonly getPublicSchemas: (region: AwsRegion) => Promise, getPrivateResources: () => Promise, getSamSchemas: () => Promise>, - private profile: string = DefaultSettings.profile.profile, - private readonly onSchemaUpdate: (region?: string, profile?: string) => void, ) { - this.privateTask = new GetPrivateSchemasTask(getPrivateResources, () => this.profile); + this.privateTask = new GetPrivateSchemasTask(getPrivateResources); this.samTask = new GetSamSchemaTask(getSamSchemas); - - this.timeout = setTimeout(() => { - // Wait before trying to call CFN APIs so that credentials have time to update - this.runPrivateTask(); - }, TenSeconds); - - this.interval = setInterval(() => { - // Keep private schemas up to date with credential changes if profile has not already ben loaded - this.runPrivateTask(); - }, OneHour); } - configure(settingsManager: ISettingsSubscriber): void { - // Clean up existing subscription if present - if (this.settingsSubscription) { - this.settingsSubscription.unsubscribe(); - } - - // Set initial settings - this.profile = settingsManager.getCurrentSettings().profile.profile; + addTask(reg: string, regionFirstCreatedMs?: number) { + const region = getRegion(reg); - // Subscribe to profile settings changes - this.settingsSubscription = settingsManager.subscribe('profile', (newProfileSettings) => { - this.onSettingsChanged(newProfileSettings.profile); - }); - } + if (!this.processedRegions.has(region)) { + this.tasks.push(new GetPublicSchemaTask(region, this.getPublicSchemas, regionFirstCreatedMs)); + this.processedRegions.add(region); + } - private onSettingsChanged(newProfile: string): void { - this.profile = newProfile; + if (!this.isRunning) { + this.runNextTask(); + } } - addTask(region: AwsRegion, regionFirstCreatedMs?: number) { - if (!this.currentRegionalTasks().has(region)) { - this.tasks.push(new GetPublicSchemaTask(region, this.getPublicSchemas, regionFirstCreatedMs)); + private runNextTask() { + const task = this.tasks.shift(); + if (!task) { + this.isRunning = false; + return; } - this.startProcessing(); + + this.isRunning = true; + task.run(this.schemas.publicSchemas) + .catch((err) => { + this.log.error(err); + this.tasks.push(task); + }) + .finally(() => { + this.isRunning = false; + this.runNextTask(); + }); } runPrivateTask() { this.privateTask - .run(this.schemas.privateSchemas, this.log) - .then(() => { - this.onSchemaUpdate(undefined, this.profile); - }) - .catch(() => {}); + .run(this.schemas.privateSchemas) + .then(() => this.schemas.invalidate()) + .catch(this.log.error); } runSamTask() { this.samTask - .run(this.schemas.samSchemas, this.log) - .then(() => { - this.onSchemaUpdate(); // No params = SAM update - }) - .catch(() => {}); - } - - public currentRegionalTasks() { - return new Set(this.tasks.map((task) => task.region)); - } - - private startProcessing() { - if (!this.isRunning && this.tasks.length > 0) { - this.isRunning = true; - this.run(); - } - } - - private run() { - const task = this.tasks.shift(); - if (task) { - task.run(this.schemas.publicSchemas, this.log) - .then(() => { - this.onSchemaUpdate(task.region); - }) - .catch(() => { - this.tasks.push(task); - }) - .finally(() => { - this.isRunning = false; - this.startProcessing(); - }); - } - } - - public close() { - // Unsubscribe from settings changes - if (this.settingsSubscription) { - this.settingsSubscription.unsubscribe(); - this.settingsSubscription = undefined; - } - - clearTimeout(this.timeout); - clearInterval(this.interval); + .run(this.schemas.samSchemas) + .then(() => this.schemas.invalidate()) + .catch(this.log.error); } } diff --git a/src/schema/PrivateSchemas.ts b/src/schema/PrivateSchemas.ts index ff1749c0..506c8123 100644 --- a/src/schema/PrivateSchemas.ts +++ b/src/schema/PrivateSchemas.ts @@ -9,6 +9,7 @@ export type PrivateSchemasType = { lastModifiedMs: number; }; +export const PrivateStoreKey = 'PrivateSchemas'; export class PrivateSchemas { static readonly V1 = 'v1'; diff --git a/src/schema/SamSchemas.ts b/src/schema/SamSchemas.ts index d0a6bfb6..9d312409 100644 --- a/src/schema/SamSchemas.ts +++ b/src/schema/SamSchemas.ts @@ -19,12 +19,7 @@ export class SamSchemas { readonly lastModifiedMs: number; readonly schemas: Map; - constructor( - version: string, - schemas: { name: string; content: string; createdMs: number }[], - firstCreatedMs: number, - lastModifiedMs: number, - ) { + constructor(version: string, schemas: SchemaFileType[], firstCreatedMs: number, lastModifiedMs: number) { this.version = version; this.firstCreatedMs = firstCreatedMs; this.lastModifiedMs = lastModifiedMs; diff --git a/src/schema/SchemaRetriever.ts b/src/schema/SchemaRetriever.ts index c00ff1e1..20e7dd41 100644 --- a/src/schema/SchemaRetriever.ts +++ b/src/schema/SchemaRetriever.ts @@ -6,11 +6,10 @@ import { LoggerFactory } from '../telemetry/LoggerFactory'; import { ScopedTelemetry } from '../telemetry/ScopedTelemetry'; import { Telemetry, Measure } from '../telemetry/TelemetryDecorator'; import { Closeable } from '../utils/Closeable'; -import { AwsRegion, getRegion } from '../utils/Region'; +import { AwsRegion } from '../utils/Region'; import { CombinedSchemas } from './CombinedSchemas'; import { GetSchemaTaskManager } from './GetSchemaTaskManager'; -import { RegionalSchemasType, SchemaFileType } from './RegionalSchemas'; -import { SamSchemasType, SamStoreKey } from './SamSchemas'; +import { SchemaFileType } from './RegionalSchemas'; import { CloudFormationResourceSchema } from './SamSchemaTransformer'; import { SchemaStore } from './SchemaStore'; @@ -20,31 +19,37 @@ export class SchemaRetriever implements SettingsConfigurable, Closeable { private settingsSubscription?: SettingsSubscription; private settings: ProfileSettings = DefaultSettings.profile; private readonly log = LoggerFactory.getLogger(SchemaRetriever); - private readonly schemaTaskManager: GetSchemaTaskManager; @Telemetry() private readonly telemetry!: ScopedTelemetry; constructor( private readonly schemaStore: SchemaStore, - private readonly getPublicSchemas: (region: AwsRegion) => Promise, + getPublicSchemas: (region: AwsRegion) => Promise, getPrivateResources: () => Promise, getSamSchemas: () => Promise>, - ) { - this.schemaTaskManager = new GetSchemaTaskManager( - this.schemaStore, - this.getPublicSchemas, + private readonly schemaTaskManager: GetSchemaTaskManager = new GetSchemaTaskManager( + schemaStore, + getPublicSchemas, getPrivateResources, getSamSchemas, - this.settings.profile, - (region, profile) => this.rebuildAffectedCombinedSchemas(region, profile), - ); + ), + ) { + this.telemetry.registerGaugeProvider('schema.public.maxAge', () => this.schemaStore.getPublicSchemasMaxAge(), { + unit: 'ms', + }); - this.telemetry.registerGaugeProvider('schema.public.maxAge', () => this.getPublicSchemaMaxAge(), { + this.telemetry.registerGaugeProvider('schema.sam.maxAge', () => this.schemaStore.getSamSchemaAge(), { unit: 'ms', }); - this.telemetry.registerGaugeProvider('schema.sam.maxAge', () => this.getSamSchemaAge(), { unit: 'ms' }); + this.getRegionalSchemasIfMissing([this.settings.region]); + } + + initialize() { + this.getRegionalSchemasIfStale(); + this.getSamSchemasIfMissingOrStale(); + this.schemaTaskManager.runPrivateTask(); } configure(settingsManager: ISettingsSubscriber): void { @@ -53,101 +58,26 @@ export class SchemaRetriever implements SettingsConfigurable, Closeable { this.settingsSubscription.unsubscribe(); } - // Set initial settings - const newSettings = settingsManager.getCurrentSettings().profile; - this.onSettingsChanged(newSettings); - - // Initialize schemas with current region - this.getRegionalSchemasIfMissing([this.settings.region]); - this.getRegionalSchemasIfStale(); - this.getSamSchemasIfMissingOrStale(); - // Subscribe to profile settings changes this.settingsSubscription = settingsManager.subscribe('profile', (newProfileSettings) => { - this.onSettingsChanged(newProfileSettings); + this.getRegionalSchemasIfMissing([newProfileSettings.region]); + this.schemaTaskManager.runPrivateTask(); + this.settings = newProfileSettings; }); } - private onSettingsChanged(newSettings: ProfileSettings): void { - const regionChanged = this.settings.region !== newSettings.region; - const profileChanged = this.settings.profile !== newSettings.profile; - - if (regionChanged || profileChanged) { - // Update private schemas when profile changes - if (profileChanged) { - this.updatePrivateSchemas(); - } - - // Ensure we have schemas for the new region - if (regionChanged) { - this.getRegionalSchemasIfMissing([newSettings.region]); - } - } - this.settings = newSettings; - } - - close(): void { - if (this.settingsSubscription) { - this.settingsSubscription.unsubscribe(); - this.settingsSubscription = undefined; - } - } - getDefault(): CombinedSchemas { return this.get(this.settings.region, this.settings.profile); } @Measure({ name: 'getSchemas' }) get(region: AwsRegion, profile: string): CombinedSchemas { - // Check if combined schemas are already cached first - const cachedCombined = this.schemaStore.get(region, profile); - if (cachedCombined) { - return cachedCombined; - } - - // Only do expensive regional check if no cached combined schemas - const regionalSchemas = this.getRegionalSchemasFromStore(region); - if (!regionalSchemas) { - this.schemaTaskManager.addTask(region); - return CombinedSchemas.from(); - } - - return this.schemaStore.put(region, profile, regionalSchemas); - } - - updatePrivateSchemas() { - this.schemaStore.invalidateCombinedSchemas(); - this.schemaTaskManager.runPrivateTask(); - } - - // Surgically rebuild affected combined schemas - @Measure({ name: 'rebuildAffectedSchemas' }) - rebuildAffectedCombinedSchemas(updatedRegion?: string, updatedProfile?: string) { - if (!updatedRegion && !updatedProfile) { - // SAM update - affects all schemas - this.schemaStore.invalidateCombinedSchemas(); - this.get(this.settings.region, this.settings.profile); - return; - } - - const keys = this.schemaStore.combinedSchemas.keys(1000); - for (const key of keys) { - const [region, profile] = key.split(':'); - if ((updatedRegion && region === updatedRegion) || (updatedProfile && profile === updatedProfile)) { - // Invalidate and rebuild this specific combined schema - this.schemaStore.combinedSchemas.remove(key).catch(this.log.error); - this.get(region as AwsRegion, profile); - } - } + return this.schemaStore.get(region, profile); } - private getRegionalSchemasFromStore(region: AwsRegion) { - return this.schemaStore.publicSchemas.get(region); - } - - private getRegionalSchemasIfMissing(preloadedRegions: ReadonlyArray) { - for (const region of preloadedRegions) { - const existingValue = this.getRegionalSchemasFromStore(region); + private getRegionalSchemasIfMissing(regions: ReadonlyArray) { + for (const region of regions) { + const existingValue = this.schemaStore.getPublicSchemas(region); if (existingValue === undefined) { this.schemaTaskManager.addTask(region); @@ -156,27 +86,26 @@ export class SchemaRetriever implements SettingsConfigurable, Closeable { } private getRegionalSchemasIfStale() { - for (const key of this.schemaStore.publicSchemas.keys(50)) { - const region = getRegion(key); - const existingValue = this.getRegionalSchemasFromStore(region); + for (const key of this.schemaStore.getPublicSchemaRegions()) { + const lastModifiedMs = this.schemaStore.getPublicSchemas(key)?.lastModifiedMs; - if (existingValue === undefined) { - this.log.error(`Something went wrong, cannot find existing region ${region}`); + if (lastModifiedMs === undefined) { + this.log.error(`Something went wrong, cannot find existing region ${key}`); return; } const now = DateTime.now(); - const lastModified = DateTime.fromMillis(existingValue.lastModifiedMs); + const lastModified = DateTime.fromMillis(lastModifiedMs); const isStale = now.diff(lastModified, 'days').days >= StaleDaysThreshold; if (isStale) { - this.schemaTaskManager.addTask(region); + this.schemaTaskManager.addTask(key); } } } private getSamSchemasIfMissingOrStale() { - const existingValue = this.schemaStore.samSchemas.get(SamStoreKey); + const existingValue = this.schemaStore.getSamSchemas(); if (existingValue === undefined) { this.schemaTaskManager.runSamTask(); @@ -192,24 +121,8 @@ export class SchemaRetriever implements SettingsConfigurable, Closeable { } } - private getPublicSchemaMaxAge(): number { - let maxAge = 0; - for (const key of this.schemaStore.publicSchemas.keys(50)) { - const region = getRegion(key); - const existingValue = this.getRegionalSchemasFromStore(region); - if (existingValue) { - const age = DateTime.now().diff(DateTime.fromMillis(existingValue.lastModifiedMs)).milliseconds; - maxAge = Math.max(maxAge, age); - } - } - return maxAge; - } - - private getSamSchemaAge(): number { - const existingValue = this.schemaStore.samSchemas.get(SamStoreKey); - if (!existingValue) { - return 0; - } - return DateTime.now().diff(DateTime.fromMillis(existingValue.lastModifiedMs)).milliseconds; + close(): void { + this.settingsSubscription?.unsubscribe(); + this.settingsSubscription = undefined; } } diff --git a/src/schema/SchemaStore.ts b/src/schema/SchemaStore.ts index b84ca32c..6ccae03f 100644 --- a/src/schema/SchemaStore.ts +++ b/src/schema/SchemaStore.ts @@ -1,39 +1,129 @@ +import { DateTime } from 'luxon'; import { DataStoreFactoryProvider, Persistence, StoreName } from '../datastore/DataStore'; import { LoggerFactory } from '../telemetry/LoggerFactory'; -import { AwsRegion } from '../utils/Region'; +import { ScopedTelemetry } from '../telemetry/ScopedTelemetry'; +import { Measure, Telemetry } from '../telemetry/TelemetryDecorator'; +import { AwsRegion, getRegion } from '../utils/Region'; import { CombinedSchemas } from './CombinedSchemas'; -import { PrivateSchemasType } from './PrivateSchemas'; +import { PrivateSchemasType, PrivateStoreKey } from './PrivateSchemas'; import { RegionalSchemasType } from './RegionalSchemas'; import { SamSchemasType, SamStoreKey } from './SamSchemas'; export class SchemaStore { + @Telemetry() + private readonly telemetry!: ScopedTelemetry; private readonly log = LoggerFactory.getLogger(SchemaStore); + private readonly createCombinedSchemas: typeof CombinedSchemas.from = (r, p, s) => { + return CombinedSchemas.from(r, p, s); + }; + public readonly publicSchemas = this.dataStoreFactory.get(StoreName.public_schemas, Persistence.local); public readonly privateSchemas = this.dataStoreFactory.get(StoreName.private_schemas, Persistence.memory); public readonly samSchemas = this.dataStoreFactory.get(StoreName.sam_schemas, Persistence.local); - public readonly combinedSchemas = this.dataStoreFactory.get(StoreName.combined_schemas, Persistence.memory); + + private regionalSchemaKey?: string; + private regionalSchemas?: RegionalSchemasType; + + private combined?: CombinedSchemas; constructor(private readonly dataStoreFactory: DataStoreFactoryProvider) {} - get(region: AwsRegion, profile: string) { - return this.combinedSchemas.get(cacheKey(region, profile)); + @Measure({ name: 'get' }) + get(region: AwsRegion, _profile: string): CombinedSchemas { + let rebuild = false; + + // 1. !this.regionalSchemas - First call ever, nothing cached yet + // 2. this.regionalSchemaKey !== region - Region changed from last call (e.g., us-east-1 -> eu-west-1) + // We track the key separately because different regions have different available resource types + // (e.g., some services only exist in certain regions), so we must reload when switching regions + if (!this.regionalSchemas || this.regionalSchemaKey !== region) { + const newSchemas = this.getPublicSchemas(region); + // Only update if schemas exist for this region - if they don't, keep the old region/schemas + // This prevents breaking the cache when requesting an unavailable region (e.g., not yet downloaded) + // Without this check, we'd set regionalSchemaKey to the new region but have no schemas, + // causing the condition above to never trigger again for that region + if (newSchemas) { + rebuild = true; + this.regionalSchemas = newSchemas; + this.regionalSchemaKey = region; + } + } + + this.telemetry.countBoolean('rebuild', rebuild); + + // Rebuild combined schemas only when necessary (region changed or no cache exists) + // Private and SAM schemas are fetched fresh each time since they can change independently + if (!this.combined || rebuild) { + this.combined = this.createCombinedSchemas( + this.regionalSchemas, + this.getPrivateSchemas(), + this.getSamSchemas(), + ); + + this.log.info( + { + Public: this.combined.regionalSchemas?.schemas.size ?? 0, + Private: this.combined.privateSchemas?.schemas.size ?? 0, + Sam: this.combined.samSchemas?.schemas.size ?? 0, + Total: this.combined.numSchemas, + }, + 'Combined schemas', + ); + } + + return this.combined; } - put(region: AwsRegion, profile: string, regionalSchemas?: RegionalSchemasType): CombinedSchemas { - const privateSchemas = this.privateSchemas.get(profile); - const samSchemas = this.samSchemas.get(SamStoreKey); + getPublicSchemas(region: string): RegionalSchemasType | undefined { + return this.publicSchemas.get(getRegion(region)); + } - const combined = CombinedSchemas.from(regionalSchemas, privateSchemas, samSchemas); - this.combinedSchemas.put(cacheKey(region, profile), combined).catch(this.log.error); - return combined; + getPublicSchemaRegions(): ReadonlyArray { + return this.publicSchemas.keys(50); } - invalidateCombinedSchemas() { - this.combinedSchemas.clear().catch(this.log.error); + getPrivateSchemas(): PrivateSchemasType | undefined { + return this.privateSchemas.get(PrivateStoreKey); } -} -function cacheKey(region: AwsRegion, profile: string) { - return `${region}:${profile}`; + getSamSchemas(): SamSchemasType | undefined { + return this.samSchemas.get(SamStoreKey); + } + + getSamSchemaAge(): number { + const existingValue = this.getSamSchemas(); + if (!existingValue) { + return 0; + } + + return DateTime.now().diff(DateTime.fromMillis(existingValue.lastModifiedMs)).toMillis(); + } + + getPublicSchemasMaxAge(): number { + const regions = this.getPublicSchemaRegions(); + if (regions.length === 0) { + return 0; + } + + let maxAge: number | undefined; + for (const key of regions) { + const lastModifiedMs = this.getPublicSchemas(key)?.lastModifiedMs; + + if (lastModifiedMs) { + const age = DateTime.now().diff(DateTime.fromMillis(lastModifiedMs)).toMillis(); + if (maxAge === undefined) { + maxAge = age; + } else { + maxAge = Math.max(maxAge, age); + } + } + } + + return maxAge ?? Number.MAX_SAFE_INTEGER; + } + + invalidate() { + this.combined = undefined; + } } diff --git a/tst/unit/schema/GetSchemaTask.test.ts b/tst/unit/schema/GetSchemaTask.test.ts index 30caa1ac..b79a714f 100644 --- a/tst/unit/schema/GetSchemaTask.test.ts +++ b/tst/unit/schema/GetSchemaTask.test.ts @@ -1,6 +1,4 @@ import { DescribeTypeOutput } from '@aws-sdk/client-cloudformation'; -import { Logger } from 'pino'; -import { StubbedInstance, stubInterface } from 'ts-sinon'; import { describe, it, expect, beforeEach, vi } from 'vitest'; import { DataStore } from '../../../src/datastore/DataStore'; import { MemoryStore } from '../../../src/datastore/MemoryStore'; @@ -9,7 +7,6 @@ import { AwsRegion } from '../../../src/utils/Region'; describe('GetSchemaTask', () => { let mockDataStore: DataStore; - let mockLogger: StubbedInstance; const mockSchemas = [ { @@ -22,11 +19,6 @@ describe('GetSchemaTask', () => { beforeEach(() => { vi.clearAllMocks(); mockDataStore = new MemoryStore('TestStore'); - mockLogger = stubInterface(); - - // Setup Sinon stubs - mockLogger.info.resolves(); - mockLogger.error.resolves(); }); describe('GetPublicSchemaTask', () => { @@ -35,7 +27,7 @@ describe('GetSchemaTask', () => { const task = new GetPublicSchemaTask(AwsRegion.US_EAST_1, mockGetSchemas, undefined); const dateNowSpy = vi.spyOn(Date, 'now').mockReturnValue(12345); - await task.run(mockDataStore, mockLogger); + await task.run(mockDataStore); expect(mockGetSchemas).toHaveBeenCalledWith(AwsRegion.US_EAST_1); @@ -50,10 +42,6 @@ describe('GetSchemaTask', () => { }), ); - expect( - mockLogger.info.calledWith(`${mockSchemas.length} public schemas retrieved for ${AwsRegion.US_EAST_1}`), - ).toBe(true); - dateNowSpy.mockRestore(); }); @@ -63,7 +51,7 @@ describe('GetSchemaTask', () => { const task = new GetPublicSchemaTask(AwsRegion.US_EAST_1, mockGetSchemas, firstCreatedMs); const dateNowSpy = vi.spyOn(Date, 'now').mockReturnValue(12345); - await task.run(mockDataStore, mockLogger); + await task.run(mockDataStore); const storedValue = mockDataStore.get(AwsRegion.US_EAST_1); expect(storedValue).toEqual( @@ -93,13 +81,7 @@ describe('GetSchemaTask', () => { // Force max attempts by setting attempts to max (task as any).attempts = GetPublicSchemaTask.MaxAttempts; - await task.run(mockDataStore, mockLogger); - - expect( - mockLogger.error.calledWith( - `Reached max attempts for retrieving schemas for ${AwsRegion.US_EAST_1} without success`, - ), - ).toBe(true); + await task.run(mockDataStore); const storedValue = mockDataStore.get(AwsRegion.US_EAST_1); expect(storedValue).toBeUndefined(); @@ -115,22 +97,21 @@ describe('GetSchemaTask', () => { } as DescribeTypeOutput, ]; - it('should retrieve and save private schemas for new profile', async () => { + it('should retrieve and save private schemas', async () => { const mockGetSchemas = vi.fn().mockResolvedValue(mockPrivateSchemas); - const mockGetProfile = vi.fn().mockReturnValue('test-profile'); - const task = new GetPrivateSchemasTask(mockGetSchemas, mockGetProfile); + const task = new GetPrivateSchemasTask(mockGetSchemas); const dateNowSpy = vi.spyOn(Date, 'now').mockReturnValue(98765); - await task.run(mockDataStore, mockLogger); + await task.run(mockDataStore); expect(mockGetSchemas).toHaveBeenCalled(); - const storedValue = mockDataStore.get('test-profile'); + const storedValue = mockDataStore.get('PrivateSchemas'); expect(storedValue).toEqual( expect.objectContaining({ version: 'v1', - identifier: 'test-profile', + identifier: 'PrivateSchemas', schemas: mockPrivateSchemas, firstCreatedMs: 98765, lastModifiedMs: 98765, @@ -140,34 +121,12 @@ describe('GetSchemaTask', () => { dateNowSpy.mockRestore(); }); - it('should skip retrieval for already processed profile', async () => { - const mockGetSchemas = vi.fn().mockResolvedValue(mockPrivateSchemas); - const mockGetProfile = vi.fn().mockReturnValue('processed-profile'); - const task = new GetPrivateSchemasTask(mockGetSchemas, mockGetProfile); - - // First run should process the profile - await task.run(mockDataStore, mockLogger); - expect(mockGetSchemas).toHaveBeenCalledTimes(1); - - // Second run should skip processing - mockGetSchemas.mockClear(); - - await task.run(mockDataStore, mockLogger); - expect(mockGetSchemas).not.toHaveBeenCalled(); - }); - it('should handle errors and rethrow', async () => { const error = new Error('Schema retrieval failed'); const mockGetSchemas = vi.fn().mockRejectedValue(error); - const mockGetProfile = vi.fn().mockReturnValue('error-profile'); - const task = new GetPrivateSchemasTask(mockGetSchemas, mockGetProfile); - - await expect(task.run(mockDataStore, mockLogger)).rejects.toThrow('Schema retrieval failed'); + const task = new GetPrivateSchemasTask(mockGetSchemas); - // Check that error was logged using Sinon stub - expect(mockLogger.error.called).toBe(true); - const errorCall = mockLogger.error.getCall(0); - expect(errorCall.lastArg).toContain('Failed to get private schemas'); + await expect(task.run(mockDataStore)).rejects.toThrow('Schema retrieval failed'); }); }); }); diff --git a/tst/unit/schema/GetSchemaTaskManager.test.ts b/tst/unit/schema/GetSchemaTaskManager.test.ts index 1e803047..10a0f5e1 100644 --- a/tst/unit/schema/GetSchemaTaskManager.test.ts +++ b/tst/unit/schema/GetSchemaTaskManager.test.ts @@ -1,197 +1,144 @@ -import { DescribeTypeOutput } from '@aws-sdk/client-cloudformation'; -import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import * as sinon from 'sinon'; +import { describe, it, expect, beforeEach, afterEach } from 'vitest'; import { MemoryDataStoreFactoryProvider } from '../../../src/datastore/DataStore'; import { GetSchemaTaskManager } from '../../../src/schema/GetSchemaTaskManager'; import { SchemaStore } from '../../../src/schema/SchemaStore'; -import { AwsRegion } from '../../../src/utils/Region'; -import { flushAllPromises } from '../../utils/Utils'; +import { getTestPrivateSchemas, samFileType, Schemas, SamSchemaFiles, schemaFileType } from '../../utils/SchemaUtils'; +import { waitFor } from '../../utils/Utils'; describe('GetSchemaTaskManager', () => { - let mockSchemaStore: SchemaStore; - let manager: GetSchemaTaskManager; - let mockGetPublicSchemas: ReturnType; - let mockGetPrivateResources: ReturnType; - let mockOnSchemaUpdate: ReturnType; - let mockGetSamSchema: ReturnType; + const timeout = 250; - beforeEach(() => { - vi.clearAllMocks(); - const dataStoreFactory = new MemoryDataStoreFactoryProvider(); - mockSchemaStore = new SchemaStore(dataStoreFactory); - - mockGetPublicSchemas = vi.fn(); - mockGetPrivateResources = vi.fn().mockResolvedValue([ - { - TypeName: 'Custom::TestResource', - Description: 'Test private resource', - } as DescribeTypeOutput, - ]); - mockGetSamSchema = vi.fn(); - - mockOnSchemaUpdate = vi.fn(); - - manager = new GetSchemaTaskManager( - mockSchemaStore, - mockGetPublicSchemas, - mockGetPrivateResources, - mockGetSamSchema, - 'default', - mockOnSchemaUpdate, - ); - }); + let schemaStore: SchemaStore; + let taskManager: GetSchemaTaskManager; + let getPublicSchemasStub: sinon.SinonStub; + let getPrivateResourcesStub: sinon.SinonStub; + let getSamSchemasStub: sinon.SinonStub; - afterEach(() => { - manager.close(); - }); - - it('should process multiple different regions sequentially', async () => { - let resolveFirst: (value: any) => void; - - mockGetPublicSchemas - .mockImplementationOnce( - () => - new Promise((resolve) => { - resolveFirst = resolve; - }), - ) - .mockImplementationOnce(() => new Promise(() => {})); - - // Add two different regions - manager.addTask(AwsRegion.US_EAST_1); - manager.addTask(AwsRegion.US_WEST_2); - - // First region starts immediately - expect(mockGetPublicSchemas).toHaveBeenCalledWith(AwsRegion.US_EAST_1); - expect(mockGetPublicSchemas).toHaveBeenCalledTimes(1); - - // Complete first task - resolveFirst!([{ name: 'first.json', content: '{}', createdMs: Date.now() }]); - await flushAllPromises(); - - // Now second region should start - expect(mockGetPublicSchemas).toHaveBeenCalledWith(AwsRegion.US_WEST_2); - expect(mockGetPublicSchemas).toHaveBeenCalledTimes(2); - }); + beforeEach(() => { + schemaStore = new SchemaStore(new MemoryDataStoreFactoryProvider()); - it('should not add duplicate regions to queue', () => { - manager.addTask(AwsRegion.US_EAST_1); - manager.addTask(AwsRegion.US_EAST_1); - manager.addTask(AwsRegion.US_EAST_1); + getPublicSchemasStub = sinon.stub().resolves(schemaFileType([Schemas.S3Bucket])); + getPrivateResourcesStub = sinon.stub().resolves(getTestPrivateSchemas()); + getSamSchemasStub = sinon.stub().resolves(samFileType([SamSchemaFiles.ServerlessFunction])); - // Should only have one task in queue - const regionalTasks = manager.currentRegionalTasks(); - expect(regionalTasks.size).toBe(1); - expect(regionalTasks.has(AwsRegion.US_EAST_1)).toBe(true); + taskManager = new GetSchemaTaskManager( + schemaStore, + getPublicSchemasStub, + getPrivateResourcesStub, + getSamSchemasStub, + ); }); - it('should track multiple different regions in queue', () => { - mockGetPublicSchemas.mockImplementation(() => new Promise(() => {})); - - // Add different regions while first is processing - manager.addTask(AwsRegion.US_EAST_1); - manager.addTask(AwsRegion.US_WEST_2); - manager.addTask(AwsRegion.EU_WEST_1); - manager.addTask(AwsRegion.US_EAST_1); // duplicate - - // Should track all unique regions - const regionalTasks = manager.currentRegionalTasks(); - expect(regionalTasks.size).toBe(3); - expect(regionalTasks.has(AwsRegion.US_EAST_1)).toBe(true); - expect(regionalTasks.has(AwsRegion.US_WEST_2)).toBe(true); - expect(regionalTasks.has(AwsRegion.EU_WEST_1)).toBe(true); + afterEach(() => { + sinon.restore(); }); - it('should handle slow promises without blocking other operations', async () => { - let resolvePromise: (value: any) => void; - - mockGetPublicSchemas.mockImplementation( - () => - new Promise((resolve) => { - resolvePromise = resolve; - }), - ); - - manager.addTask(AwsRegion.US_EAST_1); - - // Task should start but not complete - expect(mockGetPublicSchemas).toHaveBeenCalledWith(AwsRegion.US_EAST_1); - expect((manager as any).isRunning).toBe(true); - - // Should not have data in store yet - expect(mockSchemaStore.publicSchemas.get(AwsRegion.US_EAST_1)).toBeUndefined(); - - // Complete the task - resolvePromise!([{ name: 'test.json', content: '{}', createdMs: Date.now() }]); - await flushAllPromises(); + describe('addTask', () => { + it('should add and run public schema task for new region', async () => { + taskManager.addTask('us-east-1'); + taskManager.addTask('us-east-1'); + taskManager.addTask('eu-west-1'); + taskManager.addTask('eu-west-2'); + + await waitFor(() => { + expect(getPublicSchemasStub.calledWith('eu-west-2')).toBe(true); + expect(getPublicSchemasStub.calledWith('eu-west-1')).toBe(true); + expect(getPublicSchemasStub.calledWith('us-east-1')).toBe(true); + }, timeout); + }); - // Should have saved to datastore - expect(mockSchemaStore.publicSchemas.get(AwsRegion.US_EAST_1)).toBeDefined(); - expect((manager as any).isRunning).toBe(false); - }); + it('should store schemas in datastore after task completion', async () => { + taskManager.addTask('us-west-2'); - it('should handle task failures by retrying', async () => { - mockGetPublicSchemas - .mockRejectedValueOnce(new Error('Network error')) - .mockResolvedValueOnce([{ name: 'retry.json', content: '{}', createdMs: Date.now() }]); + await waitFor(() => { + const schemas = schemaStore.getPublicSchemas('us-west-2'); + expect(schemas).toBeDefined(); + expect(schemas?.region).toBe('us-west-2'); + }, timeout); + }); - manager.addTask(AwsRegion.US_EAST_1); + it('should preserve firstCreatedMs when provided', async () => { + const firstCreatedMs = Date.now() - 10000; + taskManager.addTask('ap-south-1', firstCreatedMs); + await waitFor(() => { + const schemas = schemaStore.getPublicSchemas('ap-south-1'); + expect(schemas?.firstCreatedMs).toBe(firstCreatedMs); + }, timeout); + }); - await flushAllPromises(); + it('should retry failed tasks', async () => { + getPublicSchemasStub.onFirstCall().rejects(new Error('Network error')); + getPublicSchemasStub.onSecondCall().resolves(schemaFileType([Schemas.EC2Instance])); - // Should have been called twice (original + retry) - expect(mockGetPublicSchemas).toHaveBeenCalledTimes(2); - expect(mockGetPublicSchemas).toHaveBeenCalledWith(AwsRegion.US_EAST_1); + taskManager.addTask('us-east-1'); + await waitFor(() => expect(getPublicSchemasStub.callCount).toBe(2), timeout); + }); }); - it('should handle private task with slow promises', async () => { - let resolvePromise: (value: any) => void; - - mockGetPrivateResources.mockImplementation( - () => - new Promise((resolve) => { - resolvePromise = resolve; - }), - ); + describe('runPrivateTask', () => { + it('should fetch private schemas', async () => { + taskManager.runPrivateTask(); + await waitFor(() => expect(getPrivateResourcesStub.called).toBe(true), timeout); + }); - manager.runPrivateTask(); + it('should store private schemas in datastore', async () => { + taskManager.runPrivateTask(); + await waitFor(() => { + const schemas = schemaStore.getPrivateSchemas(); + expect(schemas).toBeDefined(); + expect(schemas?.schemas.length).toBeGreaterThan(0); + }, timeout); + }); - expect(mockGetPrivateResources).toHaveBeenCalled(); + it('should invalidate combined schemas after completion', async () => { + const invalidateSpy = sinon.spy(schemaStore, 'invalidate'); - // Should not have data in store yet - expect(mockSchemaStore.privateSchemas.keys(10)).toHaveLength(0); + taskManager.runPrivateTask(); + await waitFor(() => expect(invalidateSpy.called).toBe(true), timeout); + }); - resolvePromise!([{ TypeName: 'Custom::Test', Description: 'Test' } as DescribeTypeOutput]); - await flushAllPromises(); + it('should handle errors gracefully', async () => { + getPrivateResourcesStub.rejects(new Error('Auth error')); - // Should have saved to private store - expect(mockSchemaStore.privateSchemas.keys(10).length).toBeGreaterThan(0); + taskManager.runPrivateTask(); + await waitFor(() => { + const schemas = schemaStore.getPrivateSchemas(); + expect(schemas).toBeUndefined(); + }, timeout); + }); }); - it('should call schema retriever callback after successful task completion', async () => { - mockGetPublicSchemas.mockResolvedValue([{ name: 'test.json', content: '{}', createdMs: Date.now() }]); - - manager.addTask(AwsRegion.US_EAST_1); - await flushAllPromises(); - - expect(mockOnSchemaUpdate).toHaveBeenCalledWith(AwsRegion.US_EAST_1); - }); + describe('runSamTask', () => { + it('should fetch SAM schemas', async () => { + taskManager.runSamTask(); + await waitFor(() => expect(getSamSchemasStub.called).toBe(true), timeout); + }); - it('should call schema retriever callback after private task completion', async () => { - manager.runPrivateTask(); - await flushAllPromises(); + it('should store SAM schemas in datastore', async () => { + taskManager.runSamTask(); + await waitFor(() => { + const schemas = schemaStore.getSamSchemas(); + expect(schemas).toBeDefined(); + expect(schemas?.schemas.length).toBeGreaterThan(0); + }, timeout); + }); - expect(mockOnSchemaUpdate).toHaveBeenCalledWith(undefined, 'default'); - }); + it('should invalidate combined schemas after completion', async () => { + const invalidateSpy = sinon.spy(schemaStore, 'invalidate'); - it('should call schema retriever callback after SAM task completion', async () => { - // Mock the SAM task to resolve immediately - vi.spyOn(manager as any, 'samTask', 'get').mockReturnValue({ - run: vi.fn().mockResolvedValue(undefined), + taskManager.runSamTask(); + await waitFor(() => expect(invalidateSpy.called).toBe(true), timeout); }); - manager.runSamTask(); - await flushAllPromises(); + it('should handle errors gracefully', async () => { + getSamSchemasStub.rejects(new Error('Download error')); - expect(mockOnSchemaUpdate).toHaveBeenCalledWith(); + taskManager.runSamTask(); + await waitFor(() => { + const schemas = schemaStore.getSamSchemas(); + expect(schemas).toBeUndefined(); + }, timeout); + }); }); }); diff --git a/tst/unit/schema/SchemaRetriever.test.ts b/tst/unit/schema/SchemaRetriever.test.ts index 777fe88a..dbddb47a 100644 --- a/tst/unit/schema/SchemaRetriever.test.ts +++ b/tst/unit/schema/SchemaRetriever.test.ts @@ -1,172 +1,188 @@ -import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; -import { MemoryDataStoreFactoryProvider } from '../../../src/datastore/DataStore'; +import { DateTime } from 'luxon'; +import * as sinon from 'sinon'; +import { StubbedInstance, stubInterface } from 'ts-sinon'; +import { describe, it, expect, beforeEach, afterEach } from 'vitest'; import { CombinedSchemas } from '../../../src/schema/CombinedSchemas'; +import { GetSchemaTaskManager } from '../../../src/schema/GetSchemaTaskManager'; import { RegionalSchemasType } from '../../../src/schema/RegionalSchemas'; +import { SamSchemasType } from '../../../src/schema/SamSchemas'; import { SchemaRetriever } from '../../../src/schema/SchemaRetriever'; import { SchemaStore } from '../../../src/schema/SchemaStore'; -import { Settings } from '../../../src/settings/Settings'; +import { ISettingsSubscriber } from '../../../src/settings/ISettingsSubscriber'; +import { DefaultSettings, Settings } from '../../../src/settings/Settings'; import { AwsRegion } from '../../../src/utils/Region'; -import { createMockSettingsManager } from '../../utils/MockServerComponents'; -import { getTestPrivateSchemas, samFileType } from '../../utils/SchemaUtils'; +import { PartialDataObserver } from '../../../src/utils/SubscriptionManager'; +import { getTestPrivateSchemas, samFileType, Schemas, SamSchemaFiles, schemaFileType } from '../../utils/SchemaUtils'; describe('SchemaRetriever', () => { - const key = AwsRegion.US_EAST_1; - - const mockSchemaData: RegionalSchemasType = { - version: 'v1', - region: AwsRegion.US_EAST_1, - schemas: [ - { - name: 'test-schema.json', - content: JSON.stringify({ - typeName: 'AWS::S3::Bucket', - properties: {}, - description: 'description', - primaryIdentifier: [], - additionalProperties: false, - }), - createdMs: 1622548800000, // 2021-06-01 - }, - ], - firstCreatedMs: 1622548800000, // 2021-06-01 - lastModifiedMs: 1625140800000, // 2021-07-01 - }; - - let schemaStore: SchemaStore; + const defaultRegion = DefaultSettings.profile.region; + const schemaStore = stubInterface(); + const getPublicSchemasStub = sinon.stub().resolves(schemaFileType([Schemas.S3Bucket])); + const getPrivateResourcesStub = sinon.stub().resolves(getTestPrivateSchemas()); + const getSamSchemasStub = sinon.stub().resolves(samFileType([SamSchemaFiles.ServerlessFunction])); + + let settingsManager: ISettingsSubscriber; + let taskManagerStub: StubbedInstance; let schemaRetriever: SchemaRetriever; - let mockGetPublicSchemas: ReturnType; - let mockGetPrivateResources: ReturnType; - let mockGetSam: ReturnType; beforeEach(() => { - vi.clearAllMocks(); - - // Mock Date.now to return a fixed timestamp - vi.spyOn(Date.prototype, 'getTime').mockImplementation(function (this: Date) { - // For the current date (when no date is provided to the constructor) - if (this.toString() === new Date().toString()) { - return 1625140800000; // 2021-07-01 - } - // For dates created with timestamps (like the lastModifiedMs) - return this.valueOf(); - }); + schemaStore.getPublicSchemaRegions.returns([]); + schemaStore.getPublicSchemas.returns(undefined); + schemaStore.getSamSchemas.returns(undefined); - const dataStoreFactory = new MemoryDataStoreFactoryProvider(); - schemaStore = new SchemaStore(dataStoreFactory); + getPublicSchemasStub.resetHistory(); + getPrivateResourcesStub.resetHistory(); + getSamSchemasStub.resetHistory(); - mockGetPublicSchemas = vi.fn().mockResolvedValue([]); - mockGetPrivateResources = vi.fn().mockResolvedValue(getTestPrivateSchemas()); - mockGetSam = vi.fn().mockResolvedValue(samFileType()); + settingsManager = { + subscribe: sinon.stub(), + getCurrentSettings: sinon.stub(), + }; - schemaRetriever = new SchemaRetriever(schemaStore, mockGetPublicSchemas, mockGetPrivateResources, mockGetSam); + taskManagerStub = stubInterface(); + schemaRetriever = new SchemaRetriever( + schemaStore, + getPublicSchemasStub, + getPrivateResourcesStub, + getSamSchemasStub, + taskManagerStub, + ); + + expect(taskManagerStub.addTask.calledWith(DefaultSettings.profile.region)).toBe(true); + expect(taskManagerStub.runPrivateTask.called).toBe(false); + expect(taskManagerStub.runSamTask.called).toBe(false); }); afterEach(() => { - vi.restoreAllMocks(); - schemaRetriever.close(); + sinon.restore(); }); - it('should check for missing schemas on initialization', () => { - // Need to configure to trigger initialization - const mockSettingsManager = { - getCurrentSettings: () => ({ profile: { region: AwsRegion.US_EAST_1, profile: 'default' } }), - subscribe: vi.fn().mockReturnValue({ unsubscribe: vi.fn() }), - }; - schemaRetriever.configure(mockSettingsManager as any); - - expect(mockGetPublicSchemas).toHaveBeenCalledWith(AwsRegion.US_EAST_1); + describe('constructor', () => { + it('should not add task if schemas already present', () => { + const newStore = stubInterface(); + const newTaskManagerStub = stubInterface(); + + newStore.getPublicSchemas.returns(createPublicSchema(1, defaultRegion)); + new SchemaRetriever( + newStore, + getPublicSchemasStub, + getPrivateResourcesStub, + getSamSchemasStub, + newTaskManagerStub, + ); + + expect(newTaskManagerStub.addTask.called).toBe(false); + }); }); - it('should get schema from database if available', async () => { - // Put data in the store first - await schemaStore.publicSchemas.put(AwsRegion.US_EAST_1, mockSchemaData); + describe('initialize', () => { + it('should add task for stale public schemas', () => { + const staleDate = DateTime.now().minus({ days: 8 }).toMillis(); + schemaStore.getPublicSchemaRegions.returns([defaultRegion]); + schemaStore.getPublicSchemas.returns(createPublicSchema(staleDate, defaultRegion)); - const result = schemaRetriever.get(AwsRegion.US_EAST_1, 'default'); + taskManagerStub.addTask.resetHistory(); + schemaRetriever.initialize(); - expect(result).toBeInstanceOf(CombinedSchemas); - expect(result?.numSchemas).toBeGreaterThan(0); - }); + expect(taskManagerStub.addTask.called).toBe(true); + }); - it('should return empty CombinedSchemas if schema is not in database', () => { - const result = schemaRetriever.get(AwsRegion.US_EAST_1, 'default'); + it('should not add task for fresh public schemas', () => { + const freshDate = DateTime.now().minus({ days: 3 }).toMillis(); + schemaStore.getPublicSchemaRegions.returns([defaultRegion]); + schemaStore.getPublicSchemas.returns(createPublicSchema(freshDate, defaultRegion)); - expect(result).toBeInstanceOf(CombinedSchemas); - expect(result.numSchemas).toBe(0); - expect(mockGetPublicSchemas).toHaveBeenCalledWith(AwsRegion.US_EAST_1); - }); + taskManagerStub.addTask.resetHistory(); + schemaRetriever.initialize(); - it('should check for stale schemas on initialization', async () => { - // Create a schema with a timestamp that's 8 days old (stale) - const staleTimestamp = 1625140800000 - 8 * 24 * 60 * 60 * 1000; - const staleSchemaData = { ...mockSchemaData, lastModifiedMs: staleTimestamp }; - - // Put stale data in the store - await schemaStore.publicSchemas.put(key, staleSchemaData); - - // Mock Date.now to return a fixed timestamp - vi.spyOn(Date.prototype, 'getTime').mockImplementation(function (this: Date) { - // For the current date (when no date is provided to the constructor) - if (this.toString() === new Date().toString()) { - return 1625140800000; // 2021-07-01 - } - // For dates created with timestamps (like the lastModifiedMs) - return this.valueOf(); + expect(taskManagerStub.addTask.called).toBe(false); }); - // Need to configure to trigger initialization - const mockSettingsManager = { - getCurrentSettings: () => ({ profile: { region: AwsRegion.US_EAST_1, profile: 'default' } }), - subscribe: vi.fn().mockReturnValue({ unsubscribe: vi.fn() }), - }; - schemaRetriever.configure(mockSettingsManager as any); + it('should run SAM task for missing SAM schemas', () => { + schemaStore.getSamSchemas.returns(undefined); + schemaRetriever.initialize(); - expect(mockGetPublicSchemas).toHaveBeenCalledWith(AwsRegion.US_EAST_1); - }); + expect(taskManagerStub.runSamTask.called).toBe(true); + }); - it('should get default schema using user settings', async () => { - // Put data in the store first - await schemaStore.publicSchemas.put(AwsRegion.US_EAST_1, mockSchemaData); + it('should run SAM task for stale SAM schemas', () => { + const staleDate = DateTime.now().minus({ days: 10 }).toMillis(); + schemaStore.getSamSchemas.returns(createSamSchema(staleDate)); - const result = schemaRetriever.getDefault(); - expect(result).toBeInstanceOf(CombinedSchemas); - }); + taskManagerStub.runSamTask.resetHistory(); + schemaRetriever.initialize(); - it('should update private schemas when called', () => { - schemaRetriever.updatePrivateSchemas(); - expect(mockGetPrivateResources).toHaveBeenCalled(); - }); + expect(taskManagerStub.runSamTask.called).toBe(true); + }); - it('should handle settings configuration', () => { - vi.spyOn(schemaRetriever, 'updatePrivateSchemas'); - - // Configure with new settings - const testSettings = { - profile: { - region: AwsRegion.US_WEST_2, - profile: 'new-profile', - }, - } as Settings; - const mockSettingsManager = createMockSettingsManager(testSettings); - schemaRetriever.configure(mockSettingsManager); - - // The configure method should trigger schema updates - expect(schemaRetriever.updatePrivateSchemas).toHaveBeenCalled(); - }); + it('should not run SAM task for fresh SAM schemas', () => { + const freshDate = DateTime.now().minus({ days: 1 }).toMillis(); + schemaStore.getSamSchemas.returns(createSamSchema(freshDate)); - it('should rebuild affected combined schemas', async () => { - // Put data in the store first - await schemaStore.publicSchemas.put(AwsRegion.US_EAST_1, mockSchemaData); + taskManagerStub.runSamTask.resetHistory(); + schemaRetriever.initialize(); + + expect(taskManagerStub.runSamTask.called).toBe(false); + }); + + it('should run private task', () => { + schemaRetriever.initialize(); + + expect(taskManagerStub.runPrivateTask.called).toBe(true); + }); + }); - // Get a schema to create a cached entry - schemaRetriever.get(AwsRegion.US_EAST_1, 'default'); + describe('configure', () => { + it('should add task for new region on settings change and run private task', () => { + let profileObserver: PartialDataObserver | undefined; + (settingsManager.subscribe as sinon.SinonStub).callsFake((path, observer) => { + if (path === 'profile') { + profileObserver = observer; + } + return { unsubscribe: sinon.stub() }; + }); + + schemaRetriever.configure(settingsManager); + taskManagerStub.addTask.resetHistory(); + + expect(taskManagerStub.addTask.calledWith('eu-west-1')).toBe(false); + expect(taskManagerStub.runPrivateTask.called).toBe(false); + + profileObserver?.({ region: 'eu-west-1' as AwsRegion, profile: 'test-profile' }); + expect(taskManagerStub.addTask.calledWith('eu-west-1')).toBe(true); + expect(taskManagerStub.runPrivateTask.called).toBe(true); + }); + }); - // Verify it's cached - expect(schemaStore.combinedSchemas.keys(10)).toHaveLength(1); + describe('get', () => { + it('should return combined schemas for default region', () => { + const combinedResult1 = stubInterface(); + schemaStore.get.withArgs(defaultRegion, DefaultSettings.profile.profile).returns(combinedResult1); - // Rebuild affected schemas - schemaRetriever.rebuildAffectedCombinedSchemas(AwsRegion.US_EAST_1); + const combinedResult2 = stubInterface(); + schemaStore.get.withArgs(AwsRegion.EU_WEST_1, 'anotherProfile').returns(combinedResult2); - // Should still have the schema (rebuilt) - expect(schemaStore.combinedSchemas.keys(10)).toHaveLength(1); + expect(schemaRetriever.getDefault()).toBe(combinedResult1); + expect(schemaRetriever.get(AwsRegion.EU_WEST_1, 'anotherProfile')).toBe(combinedResult2); + }); }); }); + +function createPublicSchema(date: number, region: string): RegionalSchemasType { + return { + version: '1', + region, + schemas: schemaFileType([Schemas.S3Bucket]), + firstCreatedMs: date, + lastModifiedMs: date, + }; +} + +function createSamSchema(date: number): SamSchemasType { + return { + version: '1', + schemas: schemaFileType([SamSchemaFiles.ServerlessFunction]), + firstCreatedMs: date, + lastModifiedMs: date, + }; +} diff --git a/tst/unit/schema/SchemaStore.test.ts b/tst/unit/schema/SchemaStore.test.ts new file mode 100644 index 00000000..077263de --- /dev/null +++ b/tst/unit/schema/SchemaStore.test.ts @@ -0,0 +1,364 @@ +import { DateTime } from 'luxon'; +import * as sinon from 'sinon'; +import { describe, it, expect, beforeEach, afterEach } from 'vitest'; +import { MemoryDataStoreFactoryProvider } from '../../../src/datastore/DataStore'; +import { PrivateStoreKey } from '../../../src/schema/PrivateSchemas'; +import { SamStoreKey } from '../../../src/schema/SamSchemas'; +import { SchemaStore } from '../../../src/schema/SchemaStore'; +import { AwsRegion } from '../../../src/utils/Region'; +import { getTestPrivateSchemas, Schemas, SamSchemaFiles, schemaFileType } from '../../utils/SchemaUtils'; + +describe('SchemaStore', () => { + let schemaStore: SchemaStore; + let createSpy: sinon.SinonSpy; + + beforeEach(() => { + schemaStore = new SchemaStore(new MemoryDataStoreFactoryProvider()); + createSpy = sinon.spy(schemaStore as any, 'createCombinedSchemas'); + }); + + afterEach(() => { + sinon.restore(); + }); + + describe('get', () => { + it('should return combined schemas for region', async () => { + const region = 'us-east-1' as AwsRegion; + await schemaStore.publicSchemas.put(region, { + version: 1, + region, + schemas: schemaFileType([Schemas.S3Bucket]), + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + const combined = schemaStore.get(region, 'default'); + expect(createSpy.callCount).toBe(1); + + expect(combined.regionalSchemas?.region).toBe(region); + expect(combined.regionalSchemas?.schemas.size).toBe(1); + + expect(combined.privateSchemas).toBeUndefined(); + expect(combined.samSchemas).toBeUndefined(); + }); + + it('should rebuild combined schemas when region changes', async () => { + await schemaStore.publicSchemas.put('us-east-1', { + version: 1, + region: 'us-east-1', + schemas: schemaFileType([Schemas.S3Bucket]), + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + await schemaStore.publicSchemas.put('eu-west-1', { + version: 1, + region: 'eu-west-1', + schemas: schemaFileType([Schemas.EC2Instance]), + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + const combined1 = schemaStore.get('us-east-1' as AwsRegion, 'default'); + const combined2 = schemaStore.get('eu-west-1' as AwsRegion, 'default'); + + expect(combined1.regionalSchemas?.region).toBe('us-east-1'); + expect(combined2.regionalSchemas?.region).toBe('eu-west-1'); + expect(createSpy.callCount).toBe(2); + }); + + it('should cache combined schemas for same region', async () => { + const region = 'us-east-1' as AwsRegion; + await schemaStore.publicSchemas.put(region, { + version: 1, + region, + schemas: schemaFileType([Schemas.S3Bucket]), + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + const combined1 = schemaStore.get(region, 'default'); + const combined2 = schemaStore.get(region, 'default'); + + expect(combined1).toBe(combined2); + expect(createSpy.callCount).toBe(1); + }); + + it('should include private schemas in combined result', async () => { + await schemaStore.publicSchemas.put('us-east-1', { + version: 1, + region: 'us-east-1', + schemas: schemaFileType([Schemas.S3Bucket]), + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + await schemaStore.privateSchemas.put(PrivateStoreKey, { + version: 1, + identifier: PrivateStoreKey, + schemas: getTestPrivateSchemas(), + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + const combined = schemaStore.get('us-east-1' as AwsRegion, 'default'); + expect(createSpy.callCount).toBe(1); + + expect(combined.regionalSchemas?.region).toBe('us-east-1'); + expect(combined.regionalSchemas?.schemas.size).toBe(1); + + expect(combined.privateSchemas?.schemas.size).toBe(2); + expect(combined.samSchemas).toBeUndefined(); + }); + + it('should include SAM schemas in combined result', async () => { + await schemaStore.publicSchemas.put('us-east-1', { + version: 1, + region: 'us-east-1', + schemas: schemaFileType([Schemas.S3Bucket]), + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + await schemaStore.samSchemas.put(SamStoreKey, { + version: 1, + schemas: schemaFileType([SamSchemaFiles.ServerlessFunction]), + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + const combined = schemaStore.get('us-east-1' as AwsRegion, 'default'); + + expect(createSpy.callCount).toBe(1); + + expect(combined.regionalSchemas?.region).toBe('us-east-1'); + expect(combined.regionalSchemas?.schemas.size).toBe(1); + + expect(combined.privateSchemas).toBeUndefined(); + expect(combined.samSchemas?.schemas.size).toBe(1); + }); + + it('should not rebuild when requesting unavailable region', async () => { + await schemaStore.publicSchemas.put('us-east-1', { + version: 1, + region: 'us-east-1', + schemas: schemaFileType([Schemas.S3Bucket]), + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + const combined1 = schemaStore.get('us-east-1' as AwsRegion, 'default'); + const combined2 = schemaStore.get('ap-south-1' as AwsRegion, 'default'); + + expect(combined1).toBe(combined2); + expect(combined2.regionalSchemas?.region).toBe('us-east-1'); + expect(createSpy.callCount).toBe(1); + }); + + it('should not rebuild combined schemas on repeated calls with same region', async () => { + await schemaStore.publicSchemas.put('us-east-1', { + version: 1, + region: 'us-east-1', + schemas: schemaFileType([Schemas.S3Bucket]), + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + schemaStore.get('us-east-1' as AwsRegion, 'default'); + schemaStore.get('us-east-1' as AwsRegion, 'default'); + schemaStore.get('us-east-1' as AwsRegion, 'default'); + + expect(createSpy.callCount).toBe(1); + }); + + it('should rebuild only when regional data is available', async () => { + const firstRegion = 'us-east-1'; + const secondRegion = 'us-west-2'; + let combined = schemaStore.get(firstRegion as AwsRegion, 'default'); + + expect(createSpy.callCount).toBe(1); + expect(combined.regionalSchemas).toBeUndefined(); + + await schemaStore.publicSchemas.put(firstRegion, { + version: 1, + region: firstRegion, + schemas: schemaFileType([Schemas.S3Bucket]), + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + combined = schemaStore.get(firstRegion as AwsRegion, 'default'); + expect(createSpy.callCount).toBe(2); + expect(combined.regionalSchemas).toBeDefined(); + + combined = schemaStore.get(secondRegion as AwsRegion, 'default'); + expect(createSpy.callCount).toBe(2); + expect(combined.regionalSchemas?.region).toBe(firstRegion); + + await schemaStore.publicSchemas.put(secondRegion, { + version: 1, + region: secondRegion, + schemas: schemaFileType([Schemas.S3Bucket]), + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + combined = schemaStore.get(secondRegion as AwsRegion, 'default'); + expect(createSpy.callCount).toBe(3); + expect(combined.regionalSchemas?.region).toBe(secondRegion); + }); + }); + + describe('getPublicSchemas', () => { + it('should return public schemas for region', async () => { + const region = 'us-west-2'; + const schemas = schemaFileType([Schemas.S3Bucket, Schemas.EC2Instance]); + await schemaStore.publicSchemas.put(region, { + version: 1, + region, + schemas, + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + expect(schemaStore.getPublicSchemas(region)?.schemas).toEqual(schemas); + expect(schemaStore.getPublicSchemas('ap-south-1')).toBeUndefined(); + }); + + it('should return undefined for missing region', () => { + expect(schemaStore.getPublicSchemas('us-east-1')).toBeUndefined(); + }); + }); + + describe('getPublicSchemaRegions', () => { + it('should return all stored regions', async () => { + await schemaStore.publicSchemas.put('us-east-1', { + version: 1, + region: 'us-east-1', + schemas: schemaFileType([Schemas.S3Bucket]), + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + await schemaStore.publicSchemas.put('eu-west-1', { + version: 1, + region: 'eu-west-1', + schemas: schemaFileType([Schemas.EC2Instance]), + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + expect(schemaStore.getPublicSchemaRegions()).toEqual(['us-east-1', 'eu-west-1']); + }); + + it('should return empty array when no regions stored', () => { + expect(schemaStore.getPublicSchemaRegions()).toEqual([]); + }); + }); + + describe('getPrivateSchemas', () => { + it('should return private schemas', async () => { + const privateSchemas = getTestPrivateSchemas(); + await schemaStore.privateSchemas.put(PrivateStoreKey, { + version: 1, + identifier: PrivateStoreKey, + schemas: privateSchemas, + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + expect(schemaStore.getPrivateSchemas()?.schemas.length).toBe(privateSchemas.length); + }); + + it('should return undefined when no private schemas stored', () => { + expect(schemaStore.getPrivateSchemas()).toBeUndefined(); + }); + }); + + describe('getSamSchemas', () => { + it('should return SAM schemas', async () => { + const samSchemas = schemaFileType([SamSchemaFiles.ServerlessFunction, SamSchemaFiles.ServerlessApi]); + await schemaStore.samSchemas.put(SamStoreKey, { + version: 1, + schemas: samSchemas, + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + expect(schemaStore.getSamSchemas()?.schemas.length).toBe(2); + }); + + it('should return undefined when no SAM schemas stored', () => { + expect(schemaStore.getSamSchemas()).toBeUndefined(); + }); + }); + + describe('getSamSchemaAge', () => { + it('should return age in milliseconds for existing SAM schemas', async () => { + const pastTime = DateTime.now().minus({ days: 1 }).toMillis(); + await schemaStore.samSchemas.put(SamStoreKey, { + version: 1, + schemas: schemaFileType([SamSchemaFiles.ServerlessFunction]), + firstCreatedMs: pastTime, + lastModifiedMs: pastTime, + }); + + const expectedAge = DateTime.now().toMillis() - pastTime; + expect(Math.abs(schemaStore.getSamSchemaAge() - expectedAge)).toBeLessThanOrEqual(60 * 1000); + }); + + it('should return 0 when no SAM schemas exist', () => { + expect(schemaStore.getSamSchemaAge()).toBe(0); + }); + }); + + describe('getPublicSchemasMaxAge', () => { + it('should return max age across all regions', async () => { + const oldTime = DateTime.now().minus({ days: 5 }).toMillis(); + const newTime = DateTime.now().minus({ days: 1 }).toMillis(); + + await schemaStore.publicSchemas.put('us-east-1', { + version: 1, + region: 'us-east-1', + schemas: schemaFileType([Schemas.S3Bucket]), + firstCreatedMs: oldTime, + lastModifiedMs: oldTime, + }); + + await schemaStore.publicSchemas.put('eu-west-1', { + version: 1, + region: 'eu-west-1', + schemas: schemaFileType([Schemas.EC2Instance]), + firstCreatedMs: newTime, + lastModifiedMs: newTime, + }); + + const expectedAge = DateTime.now().toMillis() - oldTime; + expect(Math.abs(schemaStore.getPublicSchemasMaxAge() - expectedAge)).toBeLessThanOrEqual(60 * 1000); + }); + + it('should return 0 when no public schemas exist', () => { + expect(schemaStore.getPublicSchemasMaxAge()).toBe(0); + }); + }); + + describe('invalidate', () => { + it('should clear cached combined schemas', async () => { + await schemaStore.publicSchemas.put('us-east-1', { + version: 1, + region: 'us-east-1', + schemas: schemaFileType([Schemas.S3Bucket]), + firstCreatedMs: Date.now(), + lastModifiedMs: Date.now(), + }); + + const combined1 = schemaStore.get('us-east-1' as AwsRegion, 'default'); + expect(createSpy.callCount).toBe(1); + + schemaStore.invalidate(); + const combined2 = schemaStore.get('us-east-1' as AwsRegion, 'default'); + expect(combined1).not.toBe(combined2); + expect(createSpy.callCount).toBe(2); + }); + }); +}); diff --git a/tst/utils/TestExtension.ts b/tst/utils/TestExtension.ts index 0957c9ca..7b2a96c2 100644 --- a/tst/utils/TestExtension.ts +++ b/tst/utils/TestExtension.ts @@ -54,7 +54,6 @@ 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'; @@ -189,6 +188,10 @@ export class TestExtension implements Closeable { } get components() { + if (this.core === undefined || this.external === undefined || this.providers === undefined) { + throw new Error('LSP server has not fully initialized yet'); + } + return { ...this.core, ...this.external, @@ -203,8 +206,8 @@ export class TestExtension implements Closeable { await WaitFor.waitFor(() => { const store = this.external.schemaStore; - const pbSchemas = store?.publicSchemas?.get(DefaultSettings.profile.region); - const samSchemas = store?.samSchemas?.get(SamStoreKey); + const pbSchemas = store?.getPublicSchemas(DefaultSettings.profile.region); + const samSchemas = store?.getSamSchemas(); if (pbSchemas === undefined || samSchemas === undefined) { throw new Error('Schemas not loaded yet'); diff --git a/tst/utils/Utils.ts b/tst/utils/Utils.ts index a177adcc..74a55bf0 100644 --- a/tst/utils/Utils.ts +++ b/tst/utils/Utils.ts @@ -4,6 +4,9 @@ export async function flushAllPromises() { await setImmediate(); } +const defaultTimeoutMs = 100; +const defaultIntervalMs = 5; + export class WaitFor { constructor( private readonly maxWaitMs: number, @@ -35,9 +38,17 @@ export class WaitFor { static async waitFor( code: () => void | Promise, - timeoutMs: number = 100, - intervalMs: number = 5, + timeoutMs: number = defaultTimeoutMs, + intervalMs: number = defaultIntervalMs, ): Promise { await new WaitFor(timeoutMs, intervalMs).wait(code); } } + +export async function waitFor( + code: () => void | Promise, + timeoutMs: number = defaultTimeoutMs, + intervalMs: number = defaultIntervalMs, +): Promise { + await WaitFor.waitFor(code, timeoutMs, intervalMs); +}