Building API Gateway WebSocket with Mock Integration

Takahiro Iwasa

Jul 7, 2024

3 min read

API Gateway WebSocket Mock Integration

AWS users can quickly build WebSocket servers using API Gateway with mock integration. This tutorial provides step-by-step instructions on creating and testing an API Gateway WebSocket setup.

For detailed documentation, refer to the official AWS guide.

Begin by creating a CloudFormation template with the following content.

• The messageId is passed through using $input.path('$.messageId') on line 68. Refer to the mapping template reference.
• Integration response keys on lines 77 and 84 must be statusCode described in AWS official documentation.
• Response templates for the message route on lines 85-87 modify responses based on the messageId value.

AWSTemplateFormatVersion: 2010-09-09
Description: API Gateway WebSocket with Mock Integration
Resources:
  ApiGatewayV2Api:
    Type: AWS::ApiGatewayV2::Api
    Properties:
      Name: api-gateway-websocket-with-mock-integration
      ProtocolType: WEBSOCKET
      RouteSelectionExpression: $request.body.action

  ApiGatewayV2Stage:
    Type: AWS::ApiGatewayV2::Stage
    Properties:
      StageName: production
      ApiId: !Ref ApiGatewayV2Api
      AutoDeploy: true

  ApiGatewayV2RouteOnConnect:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !Ref ApiGatewayV2Api
      RouteKey: $connect
      RouteResponseSelectionExpression: $default
      Target: !Sub integrations/${ApiGatewayV2IntegrationOnConnect}

  ApiGatewayV2RouteOnMessage:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !Ref ApiGatewayV2Api
      RouteKey: message
      RouteResponseSelectionExpression: $default
      Target: !Sub integrations/${ApiGatewayV2IntegrationOnMessage}

  ApiGatewayV2RouteResponseOnConnect:
    Type: AWS::ApiGatewayV2::RouteResponse
    Properties:
      ApiId: !Ref ApiGatewayV2Api
      RouteResponseKey: $default
      RouteId: !Ref ApiGatewayV2RouteOnConnect

  ApiGatewayV2RouteResponseOnMessage:
    Type: AWS::ApiGatewayV2::RouteResponse
    Properties:
      ApiId: !Ref ApiGatewayV2Api
      RouteResponseKey: $default
      RouteId: !Ref ApiGatewayV2RouteOnMessage

  ApiGatewayV2IntegrationOnConnect:
    Type: AWS::ApiGatewayV2::Integration
    Properties:
      ApiId: !Ref ApiGatewayV2Api
      ConnectionType: INTERNET
      IntegrationType: MOCK
      PassthroughBehavior: WHEN_NO_MATCH
      RequestTemplates:
        '$default': '{"statusCode": 200}'
      TimeoutInMillis: 29000
      PayloadFormatVersion: '1.0'

  ApiGatewayV2IntegrationOnMessage:
    Type: AWS::ApiGatewayV2::Integration
    Properties:
      ApiId: !Ref ApiGatewayV2Api
      ConnectionType: INTERNET
      IntegrationType: MOCK
      PassthroughBehavior: WHEN_NO_MATCH
      RequestTemplates:
        '$default': '{"statusCode": 200, "messageId": $input.path(''$.messageId'')}'
      TimeoutInMillis: 29000
      PayloadFormatVersion: '1.0'

  ApiGatewayV2IntegrationResponseOnConnect:
    Type: AWS::ApiGatewayV2::IntegrationResponse
    Properties:
      ApiId: !Ref ApiGatewayV2Api
      IntegrationId: !Ref ApiGatewayV2IntegrationOnConnect
      IntegrationResponseKey: /200/

  ApiGatewayV2IntegrationResponseOnMessage:
    Type: AWS::ApiGatewayV2::IntegrationResponse
    Properties:
      ApiId: !Ref ApiGatewayV2Api
      IntegrationId: !Ref ApiGatewayV2IntegrationOnMessage
      IntegrationResponseKey: /200/
      ResponseTemplates:
        '1': '{"message": "Hello World"}'
        '2': '{"message": "Hi!"}'
      TemplateSelectionExpression: ${request.body.messageId}

Outputs:
  ApiGatewayV2ApiEndpoint:
    Value: !GetAtt ApiGatewayV2Api.ApiEndpoint

Deploy the CloudFormation stack with the following command:

aws cloudformation deploy \
  --template-file template.yaml \
  --stack-name api-gateway-websocket-with-mock-integration

To fetch output values, use the following command:

aws cloudformation describe-stacks \
  --stack-name api-gateway-websocket-with-mock-integration \
| jq ".Stacks[0].Outputs"

Example Output:

[
  {
    "OutputKey": "ApiGatewayV2ApiEndpoint",
    "OutputValue": "wss://<id>.execute-api.<region>.amazonaws.com"
  }
]

Install wscat as a WebSocket client:

npm i wscat

Run the following command to connect to the WebSocket endpoint. The $connect route is used for new connections.

wscat -c wss://<id>.execute-api.<region>.amazonaws.com/production/

Send test messages and observe the responses based on the messageId:

> {"action": "message", "messageId": 1}
< {"message": "Hello World"}
> {"action": "message", "messageId": 2}
< {"message": "Hi!"}

Remove all provisioned AWS resources with the following command:

aws cloudformation delete-stack --stack-name api-gateway-websocket-with-mock-integration

Using mock integration in API Gateway WebSocket setups can streamline application development by enabling client-side implementation without waiting for server-side readiness. This guide simplifies the process of creating and testing WebSocket APIs.

Happy Coding! 🚀