Update testing to validate messaging

Author: kddejong

tst/e2e/Diagnostics.test.ts

@@ -1,22 +1,55 @@
+import { join } from 'path';
 import { describe, it, expect, beforeAll, afterAll, beforeEach } from 'vitest';
 import { TestExtension } from '../utils/TestExtension';
+import { WaitFor } from '../utils/Utils';
 
 describe('Diagnostic Features', () => {
-    const client = new TestExtension();
-    const diagnosticsReceived: any[] = [];
+    const client = new TestExtension({
+        initializeParams: {
+            initializationOptions: {
+                aws: {
+                    clientInfo: {
+                        extension: {
+                            name: 'Test CloudFormation Language Server',
+                            version: '1.0.0-test',
+                        },
+                        clientId: 'test-client',
+                    },
+                },
+                settings: {
+                    diagnostics: {
+                        cfnGuard: {
+                            enabled: true,
+                            rulesFile: join(__dirname, '../resources/guard/test-guard-rules.guard'),
+                            delayMs: 100,
+                            validateOnChange: true,
+                        },
+                    },
+                },
+            },
+        },
+    });
 
     beforeAll(async () => {
         await client.ready();
 
-        // Listen for diagnostic notifications
-        (client as any).clientConnection.onNotification('textDocument/publishDiagnostics', (params: any) => {
-            diagnosticsReceived.push(params);
+        // Configure guard with custom rules file
+        await client.changeConfiguration({
+            settings: {
+                diagnostics: {
+                    cfnGuard: {
+                        enabled: true,
+                        rulesFile: join(__dirname, '../resources/guard/test-guard-rules.guard'),
+                        delayMs: 100,
+                        validateOnChange: true,
+                    },
+                },
+            },
         });
     });
 
     beforeEach(async () => {
         await client.reset();
-        diagnosticsReceived.length = 0; // Clear previous diagnostics
     });
 
     afterAll(async () => {
@@ -25,59 +58,30 @@ describe('Diagnostic Features', () => {
 
     describe('Guard diagnostics while authoring', () => {
         it('should receive diagnostics during incremental typing', async () => {
-            // Start with basic template
+            // Start with basic template that should trigger our custom guard rules
             const initialTemplate = `AWSTemplateFormatVersion: '2010-09-09'
 Resources:
   MyBucket:
     Type: AWS::S3::Bucket`;
 
             const uri = await client.openYamlTemplate(initialTemplate);
 
-            // Wait for initial diagnostics
-            await new Promise((resolve) => setTimeout(resolve, 500));
-
-            // Simulate typing Properties section incrementally
-            await client.changeDocument({
-                textDocument: { uri, version: 2 },
-                contentChanges: [
-                    {
-                        range: {
-                            start: { line: 3, character: 25 },
-                            end: { line: 3, character: 25 },
-                        },
-                        text: `
-    Properties:`,
-                    },
-                ],
-            });
-
-            // Wait for diagnostics after adding BucketName
-            await new Promise((resolve) => setTimeout(resolve, 300));
+            // Wait for diagnostics from our custom guard rules
+            await WaitFor.waitFor(() => {
+                if (client.receivedDiagnostics.length === 0) {
+                    throw new Error('No diagnostics received yet');
+                }
+            }, 5000);
 
-            // Now add encryption property incrementally (fixing potential guard violation)
-            await client.changeDocument({
-                textDocument: { uri, version: 4 },
-                contentChanges: [
-                    {
-                        range: {
-                            start: { line: 5, character: 23 },
-                            end: { line: 5, character: 23 },
-                        },
-                        text: `
-      BucketEncryption:
-        ServerSideEncryptionConfiguration:
-          - ServerSideEncryptionByDefault:
-              SSEAlgorithm: AES256`,
-                    },
-                ],
-            });
+            expect(client.receivedDiagnostics.length).toBeGreaterThan(0);
 
-            // Wait for final diagnostics
-            await new Promise((resolve) => setTimeout(resolve, 500));
+            const latestDiagnostics = client.receivedDiagnostics[client.receivedDiagnostics.length - 1];
+            expect(latestDiagnostics.uri).toBe(uri);
+            expect(latestDiagnostics.diagnostics.length).toBeGreaterThan(0);
 
-            // Test validates that incremental document changes trigger diagnostic notifications
-            expect(uri).toBeDefined();
-            // In a real environment with guard enabled, we would validate diagnostic content here
+            // Verify we got our custom guard diagnostics
+            const guardDiagnostics = latestDiagnostics.diagnostics.filter((d: any) => d.source === 'cfn-guard');
+            expect(guardDiagnostics.length).toBeGreaterThan(0);
 
             await client.closeDocument({ textDocument: { uri } });
         });

tst/resources/guard/test-guard-rules.guard

@@ -0,0 +1,18 @@
+# Test Guard Rules for E2E Testing
+# Simple rules that will trigger on common CloudFormation patterns
+
+rule test_bucket_must_have_encryption {
+    Resources.*[ Type == "AWS::S3::Bucket" ] {
+        Properties.BucketEncryption exists <<
+            S3 buckets must have encryption configured
+        >>
+    }
+}
+
+rule test_bucket_name_required {
+    Resources.*[ Type == "AWS::S3::Bucket" ] {
+        Properties.BucketName exists <<
+            S3 buckets must have a BucketName specified
+        >>
+    }
+}

tst/utils/TestExtension.ts

@@ -81,6 +81,8 @@ type TestExtensionConfig = {
 export class TestExtension implements Closeable {
     private readonly awsMetadata: AwsMetadata;
     private readonly initializeParams: ExtendedInitializeParams;
+    private readonly diagnosticsReceived: any[] = [];
+    private readonly mockWorkspaceConfig: Record<string, unknown> = {};
 
     private readonly readStream = new PassThrough();
     private readonly writeStream = new PassThrough();
@@ -179,8 +181,12 @@ export class TestExtension implements Closeable {
         );
 
         // Handle workspace/configuration requests from the server
-        this.clientConnection.onRequest('workspace/configuration', () => {
-            return config.workspaceConfig ?? [{}];
+        this.clientConnection.onRequest('workspace/configuration', (params: any) => {
+            // Handle both array and single section requests
+            if (Array.isArray(params)) {
+                return params.map((item: any) => this.mockWorkspaceConfig[item.section] ?? {});
+            }
+            return [this.mockWorkspaceConfig];
         });
 
         this.serverConnection.listen();
@@ -204,6 +210,11 @@ export class TestExtension implements Closeable {
             await this.clientConnection.sendRequest(InitializeRequest.type, this.initializeParams);
             await this.clientConnection.sendNotification(InitializedNotification.type, {});
 
+            // Set up diagnostic listener
+            this.clientConnection.onNotification('textDocument/publishDiagnostics', (params: any) => {
+                this.diagnosticsReceived.push(params);
+            });
+
             await WaitFor.waitFor(() => {
                 const store = this.external.schemaStore;
                 const pbSchemas = store?.getPublicSchemas(DefaultSettings.profile.region);
@@ -229,6 +240,7 @@ export class TestExtension implements Closeable {
         this.core.awsCredentials.handleIamCredentialsDelete();
         this.core.usageTracker.clear();
         this.core.validationManager.clear();
+        this.diagnosticsReceived.length = 0;
     }
 
     async send(method: string, params: any) {
@@ -342,13 +354,21 @@ export class TestExtension implements Closeable {
     }
 
     changeConfiguration(params: DidChangeConfigurationParams) {
+        // Update mock workspace config and notify
+        if (params.settings) {
+            Object.assign(this.mockWorkspaceConfig, params.settings);
+        }
         return this.notify(DidChangeConfigurationNotification.method, params);
     }
 
     changeWorkspaceFolders(params: DidChangeWorkspaceFoldersParams) {
         return this.notify(DidChangeWorkspaceFoldersNotification.method, params);
     }
 
+    get receivedDiagnostics() {
+        return this.diagnosticsReceived;
+    }
+
     updateIamCredentials(params: UpdateCredentialsParams) {
         return this.send(IamCredentialsUpdateRequest.method, params);
     }