DEV Community

Cover image for AWS WebSocket Wonders – Part 2: Build & Test a WebSocket API
AWS Community Builders profile image Utkarsh Rastogi
Utkarsh Rastogi for AWS Community Builders

Posted on

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
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

🚀 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
Enter fullscreen mode Exit fullscreen mode

Step 2: Deploy Lambda Functions 

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

Step 3: Deploy API Gateway WebSocket API 

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

🧪 Test with Postman

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

1️⃣ Connect to WebSocket API

  1. Open Postman and create a new request.
  2. Change the request type to WebSocket (dropdown to the left of the URL bar).

  3. In the URL field, paste your WebSocket endpoint from the CloudFormation output:

    wss://<API_ID>.execute-api.<REGION>.amazonaws.com/dev

  4. Click Connect

Connect

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

2️⃣ Send a Message

  1. In the WebSocket message box, paste: 
{ "action": "sendmessage", "message": "Welcome To WebSocket Learning!" }
Enter fullscreen mode Exit fullscreen mode

2.Click Send

Send

✅ This triggers the SendMessageFunction Lambda. Check CloudWatch Logs.

3️⃣ Disconnect

  1. In Postman, click Disconnect in the WebSocket tab.
  2. ✅ This triggers the OnDisconnectFunction Lambda.
  3. 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

Top comments (0)