Utkarsh Rastogi for AWS Community Builders

Posted on Aug 3

AWS WebSocket Wonders – Part 2: Build & Test a WebSocket API

#aws #websocket #apigateway #serverless

In Part 1, we learned what WebSocket APIs are and why they’re awesome for real-time apps.

In Part 2, we’ll deploy our WebSocket API using separate CloudFormation templates for IAM, Lambda, and API Gateway — and then test it in Postman. 🚀

🧠 Why Modular Templates?

Splitting CloudFormation into multiple templates:

Keeps code clean and manageable
Lets teams work on different stacks independently
Makes redeployment faster
Improves reusability for future projects

1️⃣ IAM Role Template – `iam.yaml`

AWSTemplateFormatVersion: 2010-09-09
Description: IAM Role for WebSocket Lambdas

Parameters:
  LambdaRoleName:
    Type: String
    Default: WebSocketLambdaExecutionRole
    Description: Lambda role


Resources:
  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Ref LambdaRoleName
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: LambdaLogsPolicys
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: "*"
      Tags:
        - Key: Project
          Value: WebSocketAPI
        - Key: Owner
          Value: Utkarsh

Outputs:
  LambdaExecutionRoleArn:
    Description: ARN of the Lambda Execution Role
    Value: !GetAtt LambdaExecutionRole.Arn
    Export:
      Name: LambdaExecutionRoleArn

2️⃣ Lambda Functions Template – `lambda.yaml`

AWSTemplateFormatVersion: 2010-09-09
Description: Lambda functions for WebSocket API with Tags and CloudWatch Log Groups

Parameters:
  LambdaRoleName:
    Type: String
    Default: WebSocketLambdaExecutionRole
    Description: ARN of IAM Role for Lambdas

Resources:
  ########################
  # LAMBDA FUNCTIONS
  ########################
  OnConnectFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: OnConnectFunction
      Handler: index.lambda_handler
      Role: !Sub arn:aws:iam::${AWS::AccountId}:role/${LambdaRoleName}
      Runtime: python3.12
      Code:
        ZipFile: |
          def lambda_handler(event, context):
              print("Client connected:", event["requestContext"]["connectionId"])
              return {"statusCode": 200}
      Tags:
        - Key: Project
          Value: WebSocketAPI
        - Key: Owner
          Value: Utkarsh

  OnDisconnectFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: OnDisconnectFunction
      Handler: index.lambda_handler
      Role: !Sub arn:aws:iam::${AWS::AccountId}:role/${LambdaRoleName}
      Runtime: python3.12
      Code:
        ZipFile: |
          def lambda_handler(event, context):
              print("Client disconnected:", event["requestContext"]["connectionId"])
              return {"statusCode": 200}
      Tags:
        - Key: Project
          Value: WebSocketAPI
        - Key: Owner
          Value: Utkarsh

  SendMessageFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: SendMessageFunction
      Handler: index.lambda_handler
      Role: !Sub arn:aws:iam::${AWS::AccountId}:role/${LambdaRoleName}
      Runtime: python3.12
      Code:
        ZipFile: |
          import json
          def lambda_handler(event, context):
              body = json.loads(event["body"])
              print("Message received:", body)
              return {"statusCode": 200}
      Tags:
        - Key: Project
          Value: WebSocketAPI
        - Key: Owner
          Value: Utkarsh



  ########################
  # CLOUDWATCH LOG GROUPS
  ########################
  OnConnectLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub /aws/lambda/${OnConnectFunction}
      RetentionInDays: 14
      Tags:
        - Key: Project
          Value: WebSocketAPI
        - Key: Owner
          Value: Utkarsh

  OnDisconnectLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub /aws/lambda/${OnDisconnectFunction}
      RetentionInDays: 14
      Tags:
        - Key: Project
          Value: WebSocketAPI
        - Key: Owner
          Value: Utkarsh

  SendMessageLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub /aws/lambda/${SendMessageFunction}
      RetentionInDays: 14
      Tags:
        - Key: Project
          Value: WebSocketAPI
        - Key: Owner
          Value: Utkarsh

Outputs:
  OnConnectFunctionArn:
    Value: !GetAtt OnConnectFunction.Arn
    Export:
      Name: OnConnectFunctionArn

  OnDisconnectFunctionArn:
    Value: !GetAtt OnDisconnectFunction.Arn
    Export:
      Name: OnDisconnectFunctionArn

  SendMessageFunctionArn:
    Value: !GetAtt SendMessageFunction.Arn
    Export:
      Name: SendMessageFunctionArn

3️⃣ API Gateway WebSocket Template – `apigateway.yaml`

AWSTemplateFormatVersion: 2010-09-09
Description: API Gateway WebSocket API with Lambda integrations, Tags, and CloudWatch logging

Parameters:
  OnConnectFunctionName:
    Type: String
    Default: OnConnectFunction
    Description: OnConnect Lambda
  OnDisconnectFunctionName:
    Type: String
    Default: OnDisconnectFunction
    Description: OnDisconnect Lambda
  SendMessageFunctionName:
    Type: String
    Default: SendMessageFunction
    Description: SendMessage Lambda

Resources:
  ########################
  # API GATEWAY - WEBSOCKET API
  ########################
  ChatWebSocketAPI:
    Type: AWS::ApiGatewayV2::Api
    Properties:
      Name: ChatWebSocketAPI
      ProtocolType: WEBSOCKET
      RouteSelectionExpression: "$request.body.action"
      Tags:
        Project: WebSocketAPI
        Owner: Utkarsh

  ########################
  # INTEGRATIONS
  ########################
  ConnectIntegration:
    Type: AWS::ApiGatewayV2::Integration
    Properties:
      ApiId: !Ref ChatWebSocketAPI
      IntegrationType: AWS_PROXY
      IntegrationUri: !Sub arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:${OnConnectFunctionName}/invocations

  DisconnectIntegration:
    Type: AWS::ApiGatewayV2::Integration
    Properties:
      ApiId: !Ref ChatWebSocketAPI
      IntegrationType: AWS_PROXY
      IntegrationUri: !Sub arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:${OnDisconnectFunctionName}/invocations

  SendMessageIntegration:
    Type: AWS::ApiGatewayV2::Integration
    Properties:
      ApiId: !Ref ChatWebSocketAPI
      IntegrationType: AWS_PROXY
      IntegrationUri: !Sub arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:${SendMessageFunctionName}/invocations

  ########################
  # ROUTES
  ########################
  ConnectRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !Ref ChatWebSocketAPI
      RouteKey: $connect
      AuthorizationType: NONE
      Target: !Join ["/", ["integrations", !Ref ConnectIntegration]]

  DisconnectRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !Ref ChatWebSocketAPI
      RouteKey: $disconnect
      AuthorizationType: NONE
      Target: !Join ["/", ["integrations", !Ref DisconnectIntegration]]

  SendMessageRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !Ref ChatWebSocketAPI
      RouteKey: sendmessage
      AuthorizationType: NONE
      Target: !Join ["/", ["integrations", !Ref SendMessageIntegration]]

  ########################
  # PERMISSIONS
  ########################
  ConnectLambdaPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !Sub arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:${OnConnectFunctionName}
      Principal: apigateway.amazonaws.com

  DisconnectLambdaPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !Sub arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:${OnDisconnectFunctionName}
      Principal: apigateway.amazonaws.com

  SendMessageLambdaPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !Sub arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:${SendMessageFunctionName}
      Principal: apigateway.amazonaws.com

  ########################
  # DEPLOYMENT + STAGE
  ########################
  WebSocketDeployment:
    Type: AWS::ApiGatewayV2::Deployment
    Properties:
      ApiId: !Ref ChatWebSocketAPI
    DependsOn:
      - ConnectRoute
      - DisconnectRoute
      - SendMessageRoute

  DevStage:
    Type: AWS::ApiGatewayV2::Stage
    Properties:
      StageName: dev
      ApiId: !Ref ChatWebSocketAPI
      DeploymentId: !Ref WebSocketDeployment
      AccessLogSettings:
        DestinationArn: !GetAtt ApiGatewayLogGroup.Arn
        Format: '{"requestId":"$context.requestId","ip":"$context.identity.sourceIp","routeKey":"$context.routeKey","status":"$context.status","connectionId":"$context.connectionId"}'
      Tags:
        Project: WebSocketAPI
        Owner: Utkarsh

  ########################
  # API GATEWAY LOG GROUP
  ########################
  ApiGatewayLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub /aws/apigateway/${ChatWebSocketAPI}-dev
      RetentionInDays: 14
      Tags:
        - Key: Project
          Value: WebSocketAPI
        - Key: Owner
          Value: Utkarsh

Outputs:
  WebSocketURL:
    Description: WebSocket connection URL
    Value: !Sub wss://${ChatWebSocketAPI}.execute-api.${AWS::Region}.amazonaws.com/dev

🚀 Deploy in 3 Steps

Note: Make sure you have the AWS CLI installed and configured with appropriate credentials (aws configure) before running these commands.

Step 1: Deploy IAM Role

aws cloudformation deploy --template-file iam.yaml --stack-name WebSocket-IAM-Stack --capabilities CAPABILITY_NAMED_IAM

Step 2: Deploy Lambda Functions

aws cloudformation deploy --template-file lambda.yaml --stack-name WebSocket-Lambda-Stack --capabilities CAPABILITY_NAMED_IAM

Step 3: Deploy API Gateway WebSocket API

aws cloudformation deploy --template-file apigateway.yaml --stack-name WebSocket-APIGateway-Stack --capabilities CAPABILITY_NAMED_IAM

🧪 Test with Postman

Once the deployment is complete, you can test your WebSocket API using Postman.

1️⃣ Connect to WebSocket API

Open Postman and create a new request.
Change the request type to WebSocket (dropdown to the left of the URL bar).
In the URL field, paste your WebSocket endpoint from the CloudFormation output:
```
wss://<API_ID>.execute-api.<REGION>.amazonaws.com/dev
```
Click Connect

✅ This triggers the OnConnectFunction Lambda. Check CloudWatch Logs for a Client connected message.

2️⃣ Send a Message

In the WebSocket message box, paste:

{
  "action": "sendmessage",
  "message": "Welcome To WebSocket Learning!"
}

2.Click Send

✅ This triggers the SendMessageFunction Lambda. Check CloudWatch Logs.

3️⃣ Disconnect

In Postman, click Disconnect in the WebSocket tab.
✅ This triggers the OnDisconnectFunction Lambda.
Check CloudWatch Logs

📜 Summary

In this part, we:

Deployed an AWS WebSocket API using three separate CloudFormation templates:
- IAM Role (iam.yaml)
- Lambda Functions (lambda.yaml)
- API Gateway WebSocket (apigateway.yaml)
Connected to the WebSocket API using Postman
Triggered and verified:
- $connect → OnConnectFunction
- sendmessage → SendMessageFunction
- $disconnect → OnDisconnectFunction
Viewed CloudWatch Logs to confirm each Lambda executed successfully

With this modular approach, you now have a fully functional real-time WebSocket API ready to extend.

💡 Coming in Part 3:

We’ll integrate DynamoDB to store chat messages and update the SendMessageFunction to broadcast messages to all connected clients for a complete real-time chat experience.

👨‍💻 About Me

Hi! I'm Utkarsh, a Cloud Specialist & AWS Community Builder who loves turning complex AWS topics into fun chai-time stories ☕

👉 Explore more

DEV Community

AWS WebSocket Wonders – Part 2: Build & Test a WebSocket API

🧠 Why Modular Templates?

1️⃣ IAM Role Template – `iam.yaml`

2️⃣ Lambda Functions Template – `lambda.yaml`

3️⃣ API Gateway WebSocket Template – `apigateway.yaml`

🚀 Deploy in 3 Steps

Step 1: Deploy IAM Role

Step 2: Deploy Lambda Functions

Step 3: Deploy API Gateway WebSocket API

🧪 Test with Postman

1️⃣ Connect to WebSocket API

2️⃣ Send a Message

3️⃣ Disconnect

📜 Summary

👨‍💻 About Me

Top comments (0)

🧠 Why Modular Templates?

1️⃣ IAM Role Template – iam.yaml

2️⃣ Lambda Functions Template – lambda.yaml

3️⃣ API Gateway WebSocket Template – apigateway.yaml

🚀 Deploy in 3 Steps

Step 1: Deploy IAM Role

Step 2: Deploy Lambda Functions

Step 3: Deploy API Gateway WebSocket API

🧪 Test with Postman

1️⃣ Connect to WebSocket API

2️⃣ Send a Message

3️⃣ Disconnect

📜 Summary

👨‍💻 About Me

1️⃣ IAM Role Template – `iam.yaml`

2️⃣ Lambda Functions Template – `lambda.yaml`

3️⃣ API Gateway WebSocket Template – `apigateway.yaml`