feat(apigateway): WebSocketIntegrationResponse implementation

packages/aws-cdk-lib/aws-apigatewayv2-integrations/README.md

@@ -10,6 +10,7 @@
 - [WebSocket APIs](#websocket-apis)
   - [Lambda WebSocket Integration](#lambda-websocket-integration)
   - [AWS WebSocket Integration](#aws-websocket-integration)
+  - [Integration Responses](#integration-responses)
 
 ## HTTP APIs
 
@@ -306,4 +307,46 @@ webSocketApi.addRoute('$connect', {
 ```
 
 You can also set additional properties to change the behavior of your integration, such as `contentHandling`.
-See [Working with binary media types for WebSocket APIs](https://docs.aws.amazon.com/apigateway/latest/developerguide/websocket-api-develop-binary-media-types.html).
+See [Working with binary media types for WebSocket APIs](https://docs.aws.amazon.com/apigateway/latest/developerguide/websocket-api-develop-binary-media-types.html).
+
+
+### Integration Responses
+
+```ts
+import * as integrations from 'aws-cdk-lib/aws-apigatewayv2-integrations';
+
+const webSocketApi = new apigwv2.WebSocketApi(this, 'mywsapi');
+new apigwv2.WebSocketStage(this, 'mystage', {
+  webSocketApi,
+  stageName: 'dev',
+  autoDeploy: true,
+});
+
+declare const messageHandler: lambda.Function;
+const integration = new integrations.WebSocketLambdaIntegration(
+  'SendMessageIntegration',
+  messageHandler,
+  {
+    responses: [
+      // Default response key, will be used if no other matched
+      { responseKey: integrations.WebSocketIntegrationResponseKey.default },
+      // Success response key, will match all 2xx response HTTP status codes
+      { responseKey: integrations.WebSocketIntegrationResponseKey.success },
+      // You can also create custom response integrations for specific status codes
+      {
+        responseKey: integrations.WebSocketIntegrationResponseKey.fromStatusCode(404),
+        contentHandling: ContentHandling.CONVERT_TO_BINARY,
+        responseParameters: {
+          'method.response.header.Accept': "'application/json'",
+        },
+        templateSelectionExpression: '$default',
+        responseTemplates: {
+          'application/json': '{ "message": $context.error.message, "statusCode": 404 }',
+        },
+      },
+    ],
+  },
+);
+
+webSocketApi.addRoute('sendMessage', { integration });
+```

packages/aws-cdk-lib/aws-apigatewayv2-integrations/lib/websocket/aws.ts

@@ -1,10 +1,11 @@
 import {
-  WebSocketRouteIntegration,
   WebSocketIntegrationType,
   WebSocketRouteIntegrationConfig,
   WebSocketRouteIntegrationBindOptions,
   PassthroughBehavior,
   ContentHandling,
+  WebSocketTwoWayRouteIntegration,
+  InternalWebSocketIntegrationResponseProps,
 } from '../../../aws-apigatewayv2';
 import { IRole } from '../../../aws-iam';
 import { Duration } from '../../../core';
@@ -59,6 +60,14 @@ export interface WebSocketAwsIntegrationProps {
    */
   readonly requestTemplates?: { [contentType: string]: string };
 
+  /**
+   * Integration responses configuration
+   *
+   * @default - No response configuration provided.
+   * @see https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-websocket-api-integration-responses.html
+   */
+  readonly responses?: InternalWebSocketIntegrationResponseProps[];
+
   /**
    * The template selection expression for the integration.
    *
@@ -89,7 +98,7 @@ export interface WebSocketAwsIntegrationProps {
 /**
  * AWS WebSocket AWS Type Integration
  */
-export class WebSocketAwsIntegration extends WebSocketRouteIntegration {
+export class WebSocketAwsIntegration extends WebSocketTwoWayRouteIntegration {
   /**
    * @param id id of the underlying integration construct
    */
@@ -106,6 +115,7 @@ export class WebSocketAwsIntegration extends WebSocketRouteIntegration {
       credentialsRole: this.props.credentialsRole,
       requestParameters: this.props.requestParameters,
       requestTemplates: this.props.requestTemplates,
+      responses: this.props.responses,
       passthroughBehavior: this.props.passthroughBehavior,
       templateSelectionExpression: this.props.templateSelectionExpression,
       timeout: this.props.timeout,

packages/aws-cdk-lib/aws-apigatewayv2-integrations/lib/websocket/lambda.ts

@@ -1,9 +1,10 @@
 import {
-  WebSocketRouteIntegration,
   WebSocketIntegrationType,
   WebSocketRouteIntegrationBindOptions,
   WebSocketRouteIntegrationConfig,
   ContentHandling,
+  WebSocketTwoWayRouteIntegration,
+  InternalWebSocketIntegrationResponseProps,
 } from '../../../aws-apigatewayv2';
 import { ServicePrincipal } from '../../../aws-iam';
 import { IFunction } from '../../../aws-lambda';
@@ -28,12 +29,20 @@ export interface WebSocketLambdaIntegrationProps {
    * the route response or method response without modification.
    */
   readonly contentHandling?: ContentHandling;
+
+  /**
+   * Integration responses configuration
+   *
+   * @default - No response configuration provided.
+   * @see https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-websocket-api-integration-responses.html
+   */
+  readonly responses?: InternalWebSocketIntegrationResponseProps[];
 }
 
 /**
  * Lambda WebSocket Integration
  */
-export class WebSocketLambdaIntegration extends WebSocketRouteIntegration {
+export class WebSocketLambdaIntegration extends WebSocketTwoWayRouteIntegration {
 
   private readonly _id: string;
 
@@ -75,6 +84,7 @@ export class WebSocketLambdaIntegration extends WebSocketRouteIntegration {
       uri: integrationUri,
       timeout: this.props.timeout,
       contentHandling: this.props.contentHandling,
+      responses: this.props.responses,
     };
   }
 }

packages/aws-cdk-lib/aws-apigatewayv2-integrations/test/websocket/aws.test.ts

@@ -1,6 +1,6 @@
 import { WebSocketAwsIntegration } from './../../lib/websocket/aws';
 import { Template } from '../../../assertions';
-import { ContentHandling, PassthroughBehavior, WebSocketApi } from '../../../aws-apigatewayv2';
+import { ContentHandling, PassthroughBehavior, WebSocketApi, WebSocketIntegrationResponseKey } from '../../../aws-apigatewayv2';
 import * as iam from '../../../aws-iam';
 import { Duration, Stack } from '../../../core';
 
@@ -47,7 +47,9 @@ describe('AwsWebSocketIntegration', () => {
           passthroughBehavior: PassthroughBehavior.WHEN_NO_TEMPLATES,
           contentHandling: ContentHandling.CONVERT_TO_BINARY,
           timeout: Duration.seconds(10),
+          responses: [{ responseKey: WebSocketIntegrationResponseKey.success }],
         }),
+        returnResponse: true,
       },
     });
 
@@ -66,5 +68,10 @@ describe('AwsWebSocketIntegration', () => {
       ContentHandlingStrategy: 'CONVERT_TO_BINARY',
       TimeoutInMillis: 10000,
     });
+    Template.fromStack(stack).hasResourceProperties('AWS::ApiGatewayV2::IntegrationResponse', {
+      ApiId: { Ref: 'ApiF70053CD' },
+      IntegrationId: { Ref: 'ApidefaultRouteAwsIntegration9C4CC31F' },
+      IntegrationResponseKey: '/2\\d{2}/',
+    });
   });
 });