added to FunctionBuilder

Author: dxd1

serverless/aws/functions/websocket.py

@@ -16,9 +16,9 @@ def __init__(self, name: str, identity_source: list = None):
 class WebsocketEvent(YamlOrderedDict):
     yaml_tag = "websocket"
 
-    def __init__(self, routeKey: str = "", routeResponseSelectionExpression: str = None, authorizer: Authorizer = None):
+    def __init__(self, route: str, routeResponseSelectionExpression: str = None, authorizer: Authorizer = None):
         super().__init__()
-        self.route = routeKey
+        self.route = route
 
         if routeResponseSelectionExpression:
             self.routeResponseSelectionExpression = routeResponseSelectionExpression
@@ -35,7 +35,7 @@ def __init__(
         service,
         name,
         description,
-        routeKey="$default",
+        route="$default",
         routeResponseSelectionExpression=None,
         authorizer=None,
         handler=None,
@@ -48,4 +48,4 @@ def __init__(
         super().__init__(
             service, name, description, handler, timeout, layers, use_dlq=use_dlq, use_async_dlq=use_async_dlq, **kwargs
         )
-        self.trigger(WebsocketEvent(routeKey, routeResponseSelectionExpression, authorizer))
+        self.trigger(WebsocketEvent(route, routeResponseSelectionExpression, authorizer))

serverless/aws/provider.py

@@ -7,6 +7,7 @@
 from serverless.aws.functions.kinesis import KinesisFunction
 from serverless.aws.functions.s3 import S3Function
 from serverless.aws.functions.sqs import SQSFunction
+from serverless.aws.functions.websocket import WebsocketFunction
 from serverless.aws.iam import ServicePolicyBuilder
 from serverless.service.environment import Environment
 from serverless.service.types import Provider as BaseProvider
@@ -188,6 +189,12 @@ def sqs(self, name, description, arn, **kwargs):
 
         return fn
 
+    def websocket(self, name, description, route, **kwargs):
+        fn = WebsocketFunction(self.service, name, description, route, **kwargs)
+        self.service.functions.add(fn)
+
+        return fn
+
 
 class Provider(BaseProvider, yaml.YAMLObject):
     yaml_tag = "!Provider"