From c042c49c79769af9c73c1965805f32d6e833e6cd Mon Sep 17 00:00:00 2001 From: gemammercado Date: Tue, 21 Apr 2026 16:02:06 -0400 Subject: [PATCH] add v2 relatedResources handler --- src/handlers/RelatedResourcesHandler.ts | 28 +++ src/protocol/LspRelatedResourcesHandlers.ts | 7 +- src/protocol/RelatedResourcesProtocol.ts | 6 +- src/server/CfnServer.ts | 7 + .../handlers/RelatedResourcesHandler.test.ts | 165 +++++++++++++++++- .../LspRelatedResourcesHandlers.test.ts | 11 +- tst/unit/server/CfnServer.test.ts | 1 + 7 files changed, 217 insertions(+), 8 deletions(-) diff --git a/src/handlers/RelatedResourcesHandler.ts b/src/handlers/RelatedResourcesHandler.ts index 0ab4705..187104f 100644 --- a/src/handlers/RelatedResourcesHandler.ts +++ b/src/handlers/RelatedResourcesHandler.ts @@ -20,6 +20,34 @@ import { export function getAuthoredResourceTypesHandler( components: ServerComponents, +): RequestHandler { + return (rawParams) => { + try { + const templateUri = parseWithPrettyError(parseTemplateUriParams, rawParams); + const syntaxTree = components.syntaxTreeManager.getSyntaxTree(templateUri); + if (syntaxTree) { + const resourcesMap = getEntityMap(syntaxTree, TopLevelSection.Resources); + if (resourcesMap) { + const resourceTypes = [...resourcesMap.values()] + .map((context) => { + const resource = context.entity as Resource; + return resource?.Type; + }) + .filter((type): type is string => type !== undefined && type !== null); + + return [...new Set(resourceTypes)]; + } + } + + return []; + } catch (error) { + handleLspError(error, 'Failed to get authored resource types'); + } + }; +} + +export function getAuthoredResourceTypesHandlerV2( + components: ServerComponents, ): RequestHandler { return (rawParams) => { try { diff --git a/src/protocol/LspRelatedResourcesHandlers.ts b/src/protocol/LspRelatedResourcesHandlers.ts index 1a323c0..f21bde1 100644 --- a/src/protocol/LspRelatedResourcesHandlers.ts +++ b/src/protocol/LspRelatedResourcesHandlers.ts @@ -2,6 +2,7 @@ import { Connection, RequestHandler } from 'vscode-languageserver'; import { AuthoredResource, GetAuthoredResourceTypesRequest, + GetAuthoredResourceTypesRequestV2, GetRelatedResourceTypesParams, GetRelatedResourceTypesRequest, InsertRelatedResourcesParams, @@ -13,10 +14,14 @@ import { export class LspRelatedResourcesHandlers { constructor(private readonly connection: Connection) {} - onGetAuthoredResourceTypes(handler: RequestHandler) { + onGetAuthoredResourceTypes(handler: RequestHandler) { this.connection.onRequest(GetAuthoredResourceTypesRequest.method, handler); } + onGetAuthoredResourceTypesV2(handler: RequestHandler) { + this.connection.onRequest(GetAuthoredResourceTypesRequestV2.method, handler); + } + onGetRelatedResourceTypes(handler: RequestHandler) { this.connection.onRequest(GetRelatedResourceTypesRequest.method, handler); } diff --git a/src/protocol/RelatedResourcesProtocol.ts b/src/protocol/RelatedResourcesProtocol.ts index 64d7ea8..7932933 100644 --- a/src/protocol/RelatedResourcesProtocol.ts +++ b/src/protocol/RelatedResourcesProtocol.ts @@ -25,10 +25,14 @@ export interface RelatedResourcesCodeAction extends CodeAction { }; } -export const GetAuthoredResourceTypesRequest = new RequestType( +export const GetAuthoredResourceTypesRequest = new RequestType( 'aws/cfn/template/resources/authored', ); +export const GetAuthoredResourceTypesRequestV2 = new RequestType( + 'aws/cfn/template/resources/authored/v2', +); + export const GetRelatedResourceTypesRequest = new RequestType( 'aws/cfn/template/resources/related', ); diff --git a/src/server/CfnServer.ts b/src/server/CfnServer.ts index 7cc7fba..a367797 100644 --- a/src/server/CfnServer.ts +++ b/src/server/CfnServer.ts @@ -13,6 +13,7 @@ import { hoverHandler } from '../handlers/HoverHandler'; import { initializedHandler } from '../handlers/Initialize'; import { getAuthoredResourceTypesHandler, + getAuthoredResourceTypesHandlerV2, getRelatedResourceTypesHandler, insertRelatedResourcesHandler, } from '../handlers/RelatedResourcesHandler'; @@ -246,6 +247,12 @@ export class CfnServer { this.lsp.relatedResourcesHandlers.onGetAuthoredResourceTypes( withTelemetryContext('Related.Resources.Get.Authored', getAuthoredResourceTypesHandler(this.components)), ); + this.lsp.relatedResourcesHandlers.onGetAuthoredResourceTypesV2( + withTelemetryContext( + 'Related.Resources.Get.Authored.V2', + getAuthoredResourceTypesHandlerV2(this.components), + ), + ); this.lsp.relatedResourcesHandlers.onGetRelatedResourceTypes( withTelemetryContext('Related.Resources.Get.Related', getRelatedResourceTypesHandler(this.components)), ); diff --git a/tst/unit/handlers/RelatedResourcesHandler.test.ts b/tst/unit/handlers/RelatedResourcesHandler.test.ts index 374c73c..bac463b 100644 --- a/tst/unit/handlers/RelatedResourcesHandler.test.ts +++ b/tst/unit/handlers/RelatedResourcesHandler.test.ts @@ -3,6 +3,7 @@ import { CancellationToken } from 'vscode-languageserver'; import { getEntityMap } from '../../../src/context/SectionContextBuilder'; import { getAuthoredResourceTypesHandler, + getAuthoredResourceTypesHandlerV2, getRelatedResourceTypesHandler, insertRelatedResourcesHandler, } from '../../../src/handlers/RelatedResourcesHandler'; @@ -37,9 +38,89 @@ describe('RelatedResourcesHandler', () => { }); describe('getAuthoredResourceTypesHandler', () => { - it('should return authored resources with logical IDs and types', () => { + it('should return resource type strings', () => { + const handler = getAuthoredResourceTypesHandler(mockComponents); + const templateUri = 'file:///test/template.yaml'; + + syntaxTreeManager.getSyntaxTree.withArgs(templateUri).returns({} as any); + mockGetEntityMap.mockReturnValue( + new Map([ + ['Bucket1', { entity: { Type: 'AWS::S3::Bucket' } }], + ['Function1', { entity: { Type: 'AWS::Lambda::Function' } }], + ]), + ); + + const result = handler(templateUri, mockToken); + + expect(result).toEqual(['AWS::S3::Bucket', 'AWS::Lambda::Function']); + }); + + it('should return empty array when no syntax tree found', () => { + const handler = getAuthoredResourceTypesHandler(mockComponents); + const templateUri = 'file:///test/template.yaml'; + + syntaxTreeManager.getSyntaxTree.withArgs(templateUri).returns(undefined); + + const result = handler(templateUri, mockToken); + + expect(result).toEqual([]); + }); + + it('should return empty array when no resources found', () => { + const handler = getAuthoredResourceTypesHandler(mockComponents); + const templateUri = 'file:///test/template.yaml'; + + syntaxTreeManager.getSyntaxTree.withArgs(templateUri).returns({} as any); + mockGetEntityMap.mockReturnValue(undefined); + + const result = handler(templateUri, mockToken); + + expect(result).toEqual([]); + }); + + it('should filter out undefined and null resource types', () => { + const handler = getAuthoredResourceTypesHandler(mockComponents); + const templateUri = 'file:///test/template.yaml'; + + const mockResourceContext1 = { + entity: { Type: 'AWS::S3::Bucket' }, + }; + const mockResourceContext2 = { + entity: { Type: undefined as any }, + }; + const mockResourceContext3 = { + entity: { Type: null as any }, + }; + + syntaxTreeManager.getSyntaxTree.withArgs(templateUri).returns({} as any); + mockGetEntityMap.mockReturnValue( + new Map([ + ['Bucket1', mockResourceContext1], + ['Resource2', mockResourceContext2], + ['Resource3', mockResourceContext3], + ]) as any, + ); + + const result = handler(templateUri, mockToken); + + expect(result).toEqual(['AWS::S3::Bucket']); + }); + + it('should handle errors and rethrow them', () => { const handler = getAuthoredResourceTypesHandler(mockComponents); const templateUri = 'file:///test/template.yaml'; + const error = new Error('Syntax tree error'); + + syntaxTreeManager.getSyntaxTree.withArgs(templateUri).throws(error); + + expect(() => handler(templateUri, mockToken)).toThrow('Syntax tree error'); + }); + }); + + describe('getAuthoredResourceTypesHandlerV2', () => { + it('should return authored resources with logical IDs and types', () => { + const handler = getAuthoredResourceTypesHandlerV2(mockComponents); + const templateUri = 'file:///test/template.yaml'; const mockResourceContext1 = { entity: { Type: 'AWS::S3::Bucket' }, @@ -71,7 +152,7 @@ describe('RelatedResourcesHandler', () => { }); it('should return empty array when no syntax tree found', () => { - const handler = getAuthoredResourceTypesHandler(mockComponents); + const handler = getAuthoredResourceTypesHandlerV2(mockComponents); const templateUri = 'file:///test/template.yaml'; syntaxTreeManager.getSyntaxTree.withArgs(templateUri).returns(undefined); @@ -82,7 +163,7 @@ describe('RelatedResourcesHandler', () => { }); it('should return empty array when no resources found', () => { - const handler = getAuthoredResourceTypesHandler(mockComponents); + const handler = getAuthoredResourceTypesHandlerV2(mockComponents); const templateUri = 'file:///test/template.yaml'; syntaxTreeManager.getSyntaxTree.withArgs(templateUri).returns({} as any); @@ -94,7 +175,7 @@ describe('RelatedResourcesHandler', () => { }); it('should filter out undefined and null resource types', () => { - const handler = getAuthoredResourceTypesHandler(mockComponents); + const handler = getAuthoredResourceTypesHandlerV2(mockComponents); const templateUri = 'file:///test/template.yaml'; const mockResourceContext1 = { @@ -122,7 +203,7 @@ describe('RelatedResourcesHandler', () => { }); it('should handle errors and rethrow them', () => { - const handler = getAuthoredResourceTypesHandler(mockComponents); + const handler = getAuthoredResourceTypesHandlerV2(mockComponents); const templateUri = 'file:///test/template.yaml'; const error = new Error('Syntax tree error'); @@ -133,6 +214,46 @@ describe('RelatedResourcesHandler', () => { }); describe('getRelatedResourceTypesHandler', () => { + it('should return related resource types for a given resource type', () => { + const handler = getRelatedResourceTypesHandler(mockComponents); + const params = { parentResourceType: 'AWS::S3::Bucket' }; + + const relatedTypes = new Set(['AWS::Lambda::Function', 'AWS::IAM::Role']); + relationshipSchemaService.getAllRelatedResourceTypes.withArgs('AWS::S3::Bucket').returns(relatedTypes); + + relationshipSchemaService.getRelationshipsForResourceType.withArgs('AWS::Lambda::Function').returns({ + resourceType: 'AWS::Lambda::Function', + relationships: [ + { + property: 'BucketName', + relatedResourceTypes: [{ typeName: 'AWS::S3::Bucket', attribute: '/properties/BucketName' }], + }, + ], + }); + + relationshipSchemaService.getRelationshipsForResourceType.withArgs('AWS::IAM::Role').returns({ + resourceType: 'AWS::IAM::Role', + relationships: [ + { + property: 'BucketArn', + relatedResourceTypes: [{ typeName: 'AWS::S3::Bucket', attribute: '/properties/Arn' }], + }, + ], + }); + + mockComponents.schemaRetriever.getDefault.returns({ + schemas: new Map([ + ['AWS::Lambda::Function', { properties: { BucketName: { type: 'string' } } }], + ['AWS::IAM::Role', { properties: { BucketArn: { type: 'string' } } }], + ]), + } as any); + + const result = handler(params, mockToken); + + expect(result).toEqual(['AWS::Lambda::Function', 'AWS::IAM::Role']); + expect(relationshipSchemaService.getAllRelatedResourceTypes.calledWith('AWS::S3::Bucket')).toBe(true); + }); + it('should return related resource types that have exactly one populatable relationship', () => { const handler = getRelatedResourceTypesHandler(mockComponents); const params = { parentResourceType: 'AWS::S3::Bucket' }; @@ -343,6 +464,40 @@ describe('RelatedResourcesHandler', () => { }); describe('insertRelatedResourcesHandler', () => { + it('should insert related resources and return code action', () => { + const handler = insertRelatedResourcesHandler(mockComponents); + const params = { + templateUri: 'file:///test/template.yaml', + relatedResourceTypes: ['AWS::Lambda::Function', 'AWS::IAM::Role'], + parentResourceType: 'AWS::S3::Bucket', + }; + + const mockCodeAction = { + title: 'Insert 2 related resources', + kind: 'refactor', + edit: { + changes: { + 'file:///test/template.yaml': [], + }, + }, + }; + + mockComponents.relatedResourcesSnippetProvider.insertRelatedResources + .withArgs('file:///test/template.yaml', ['AWS::Lambda::Function', 'AWS::IAM::Role'], 'AWS::S3::Bucket') + .returns(mockCodeAction); + + const result = handler(params, mockToken); + + expect(result).toEqual(mockCodeAction); + expect( + mockComponents.relatedResourcesSnippetProvider.insertRelatedResources.calledWith( + 'file:///test/template.yaml', + ['AWS::Lambda::Function', 'AWS::IAM::Role'], + 'AWS::S3::Bucket', + ), + ).toBe(true); + }); + it('should insert related resources and return code action without parentLogicalId', () => { const handler = insertRelatedResourcesHandler(mockComponents); const params = { diff --git a/tst/unit/protocol/LspRelatedResourcesHandlers.test.ts b/tst/unit/protocol/LspRelatedResourcesHandlers.test.ts index 18e34d5..037a623 100644 --- a/tst/unit/protocol/LspRelatedResourcesHandlers.test.ts +++ b/tst/unit/protocol/LspRelatedResourcesHandlers.test.ts @@ -5,6 +5,7 @@ import { LspRelatedResourcesHandlers } from '../../../src/protocol/LspRelatedRes import { AuthoredResource, GetAuthoredResourceTypesRequest, + GetAuthoredResourceTypesRequestV2, GetRelatedResourceTypesParams, GetRelatedResourceTypesRequest, InsertRelatedResourcesParams, @@ -23,13 +24,21 @@ describe('LspRelatedResourcesHandlers', () => { }); it('should register onGetAuthoredResourceTypes handler', () => { - const mockHandler: RequestHandler = vi.fn(); + const mockHandler: RequestHandler = vi.fn(); relatedResourcesHandlers.onGetAuthoredResourceTypes(mockHandler); expect(connection.onRequest.calledWith(GetAuthoredResourceTypesRequest.method)).toBe(true); }); + it('should register onGetAuthoredResourceTypesV2 handler', () => { + const mockHandler: RequestHandler = vi.fn(); + + relatedResourcesHandlers.onGetAuthoredResourceTypesV2(mockHandler); + + expect(connection.onRequest.calledWith(GetAuthoredResourceTypesRequestV2.method)).toBe(true); + }); + it('should register onGetRelatedResourceTypes handler', () => { const mockHandler: RequestHandler = vi.fn(); diff --git a/tst/unit/server/CfnServer.test.ts b/tst/unit/server/CfnServer.test.ts index eb2933d..73a770d 100644 --- a/tst/unit/server/CfnServer.test.ts +++ b/tst/unit/server/CfnServer.test.ts @@ -48,6 +48,7 @@ describe('CfnServer', () => { expect(mockFeatures.stackHandlers.onListStackResources.calledOnce).toBe(true); expect(mockFeatures.relatedResourcesHandlers.onGetAuthoredResourceTypes.calledOnce).toBe(true); + expect(mockFeatures.relatedResourcesHandlers.onGetAuthoredResourceTypesV2.calledOnce).toBe(true); expect(mockFeatures.relatedResourcesHandlers.onGetRelatedResourceTypes.calledOnce).toBe(true); expect(mockFeatures.relatedResourcesHandlers.onInsertRelatedResources.calledOnce).toBe(true); });