allow state insertion into empty documents

Author: Zee2413

src/handlers/ResourceHandler.ts

@@ -3,6 +3,7 @@ import { ResponseError, ServerRequestHandler } from 'vscode-languageserver';
 import { RequestHandler } from 'vscode-languageserver/node';
 import { TopLevelSection } from '../context/ContextType';
 import { getEntityMap } from '../context/SectionContextBuilder';
+import { CloudFormationFileType } from '../document/Document';
 import { parseResourceTypeName } from '../resourceState/ResourceStateParser';
 import {
     ResourceTypesResult,
@@ -82,8 +83,18 @@ export function importResourceStateHandler(
 ): RequestHandler<ResourceStateParams, ResourceStateResult, void> {
     return async (params: ResourceStateParams): Promise<ResourceStateResult> => {
         components.usageTracker.track(EventType.DidImportResources);
-        if (!components.documentManager.get(params.textDocument.uri)?.isTemplate()) {
-            throw new ResponseError(400, 'Must open CloudFormation template to import or clone resource state');
+        const doc = components.documentManager.get(params.textDocument.uri);
+        if (!doc) {
+            const msg = `${params.purpose} failed: ${params.textDocument.uri} not found`;
+            log.error(msg);
+            throw new ResponseError(500, msg); // all open TextDocuments should be registered by protocol
+        }
+
+        if (!doc.isTemplate() && doc.cfnFileType !== CloudFormationFileType.Empty) {
+            throw new ResponseError(
+                400,
+                `${params.purpose} failed: ${params.textDocument.uri} is not a valid CloudFormation template`,
+            );
         }
         return await components.resourceStateImporter.importResourceState(params);
     };

tst/unit/handlers/ResourceHandler.test.ts

@@ -1,8 +1,8 @@
+import { stubInterface } from 'ts-sinon';
 import { describe, it, expect, beforeEach, vi } from 'vitest';
 import { CancellationToken, RequestHandler } from 'vscode-languageserver-protocol';
-import { TextDocument } from 'vscode-languageserver-textdocument';
 import { getEntityMap } from '../../../src/context/SectionContextBuilder';
-import { Document } from '../../../src/document/Document';
+import { CloudFormationFileType, Document } from '../../../src/document/Document';
 import {
     getManagedResourceStackTemplateHandler,
     importResourceStateHandler,
@@ -168,33 +168,39 @@ describe('ResourceHandler - removeResourceTypeHandler', () => {
 describe('ResourceHandler - importResourceStateHandler', () => {
     let mockComponents: ReturnType<typeof createMockComponents>;
     let handler: RequestHandler<ResourceStateParams, ResourceStateResult, void>;
+    const params = {
+        textDocument: { uri: 'docUri' },
+        resourceSelections: [
+            {
+                resourceType: 'AWS::S3::Bucket',
+                resourceIdentifiers: ['bucket1234'],
+            },
+        ],
+        purpose: ResourceStatePurpose.IMPORT,
+    };
 
     beforeEach(() => {
         vi.clearAllMocks();
         mockComponents = createMockComponents();
         handler = importResourceStateHandler(mockComponents);
     });
 
-    it('should throw error if template file is not open', async () => {
-        mockComponents.documentManager.get.returns(
-            new Document(TextDocument.create('sample.yaml', 'yaml', 1, 'Hello:')),
-        );
-        mockComponents.documentManager.isTemplate.resolves(false);
+    it('should throw error if document not found', async () => {
+        mockComponents.documentManager.get.returns(undefined);
 
         await expect(async () => {
-            await handler(
-                {
-                    textDocument: { uri: 'docUri' },
-                    resourceSelections: [
-                        {
-                            resourceType: 'AWS::S3::Bucket',
-                            resourceIdentifiers: ['bucket1234'],
-                        },
-                    ],
-                    purpose: ResourceStatePurpose.IMPORT,
-                },
-                CancellationToken.None,
-            );
-        }).rejects.toThrow('Must open CloudFormation template to import or clone resource state');
+            await handler(params, CancellationToken.None);
+        }).rejects.toThrow('Import failed: docUri not found');
+    });
+
+    it('should throw error if document is not a valid CloudFormation template', async () => {
+        const mockDoc = stubInterface<Document>();
+        mockDoc.isTemplate.returns(false);
+        Object.defineProperty(mockDoc, 'cfnFileType', { value: CloudFormationFileType.Other });
+        mockComponents.documentManager.get.returns(mockDoc);
+
+        await expect(async () => {
+            await handler(params, CancellationToken.None);
+        }).rejects.toThrow('Import failed: docUri is not a valid CloudFormation template');
     });
 });