Skip to content

Commit df4a42e

Browse files
fix: fix websockets not being registered (#140)
1 parent f1446ad commit df4a42e

File tree

1 file changed

+15
-20
lines changed

1 file changed

+15
-20
lines changed

nitric/resources/websockets.py

Lines changed: 15 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -72,24 +72,15 @@ async def send(self, socket: str, connection_id: str, data: bytes):
7272
raise exception_from_grpc_error(grpc_err) from grpc_err
7373

7474

75-
class WebsocketWorkerOptions:
76-
"""Options for websocket workers."""
77-
78-
def __init__(self, socket_name: str, event_type: Literal["connect", "disconnect", "message"]):
79-
"""Construct new websocket worker options."""
80-
self.socket_name = socket_name
81-
self.event_type = WebsocketWorkerOptions._to_grpc_event_type(event_type)
82-
83-
@staticmethod
84-
def _to_grpc_event_type(event_type: Literal["connect", "disconnect", "message"]) -> WebsocketEventType:
85-
if event_type == "connect":
86-
return WebsocketEventType.Connect
87-
elif event_type == "disconnect":
88-
return WebsocketEventType.Disconnect
89-
elif event_type == "message":
90-
return WebsocketEventType.Message
91-
else:
92-
raise ValueError(f"Event type {event_type} is unsupported")
75+
def _to_grpc_event_type(event_type: Literal["connect", "disconnect", "message"]) -> WebsocketEventType:
76+
if event_type == "connect":
77+
return WebsocketEventType.Connect
78+
elif event_type == "disconnect":
79+
return WebsocketEventType.Disconnect
80+
elif event_type == "message":
81+
return WebsocketEventType.Message
82+
else:
83+
raise ValueError(f"Event type {event_type} is unsupported")
9384

9485

9586
def _to_resource(b: Websocket) -> ResourceIdentifier:
@@ -194,11 +185,15 @@ class WebsocketWorker(FunctionServer):
194185
_registration_request: RegistrationRequest
195186
_responses: AsyncNotifierList[ClientMessage]
196187

197-
def __init__(self, socket_name: str, event_type: WebsocketEventType, handler: WebsocketHandler):
188+
def __init__(
189+
self, socket_name: str, event_type: Literal["connect", "disconnect", "message"], handler: WebsocketHandler
190+
):
198191
"""Construct a new WebsocketHandler."""
199192
self._handler = handler
200193
self._responses = AsyncNotifierList()
201-
self._registration_request = RegistrationRequest(socket_name=socket_name, event_type=event_type)
194+
self._registration_request = RegistrationRequest(
195+
socket_name=socket_name, event_type=_to_grpc_event_type(event_type)
196+
)
202197

203198
Nitric._register_worker(self)
204199

0 commit comments

Comments
 (0)