simple-websocket-1.1.0/0000775000175000017500000000000014702053632011712 5ustar simple-websocket-1.1.0/src/0000775000175000017500000000000014702053632012501 5ustar simple-websocket-1.1.0/src/simple_websocket/0000775000175000017500000000000014702053632016040 5ustar simple-websocket-1.1.0/src/simple_websocket/ws.py0000664000175000017500000005440514702052374017055 0ustar import selectors import socket import ssl from time import time from urllib.parse import urlsplit from wsproto import ConnectionType, WSConnection from wsproto.events import ( AcceptConnection, RejectConnection, CloseConnection, Message, Request, Ping, Pong, TextMessage, BytesMessage, ) from wsproto.extensions import PerMessageDeflate from wsproto.frame_protocol import CloseReason from wsproto.utilities import LocalProtocolError from .errors import ConnectionError, ConnectionClosed class Base: def __init__(self, sock=None, connection_type=None, receive_bytes=4096, ping_interval=None, max_message_size=None, thread_class=None, event_class=None, selector_class=None): #: The name of the subprotocol chosen for the WebSocket connection. self.subprotocol = None self.sock = sock self.receive_bytes = receive_bytes self.ping_interval = ping_interval self.max_message_size = max_message_size self.pong_received = True self.input_buffer = [] self.incoming_message = None self.incoming_message_len = 0 self.connected = False self.is_server = (connection_type == ConnectionType.SERVER) self.close_reason = CloseReason.NO_STATUS_RCVD self.close_message = None if thread_class is None: import threading thread_class = threading.Thread if event_class is None: # pragma: no branch import threading event_class = threading.Event if selector_class is None: selector_class = selectors.DefaultSelector self.selector_class = selector_class self.event = event_class() self.ws = WSConnection(connection_type) self.handshake() if not self.connected: # pragma: no cover raise ConnectionError() self.thread = thread_class(target=self._thread) self.thread.name = self.thread.name.replace( '(_thread)', '(simple_websocket.Base._thread)') self.thread.start() def handshake(self): # pragma: no cover # to be implemented by subclasses pass def send(self, data): """Send data over the WebSocket connection. :param data: The data to send. If ``data`` is of type ``bytes``, then a binary message is sent. Else, the message is sent in text format. """ if not self.connected: raise ConnectionClosed(self.close_reason, self.close_message) if isinstance(data, bytes): out_data = self.ws.send(Message(data=data)) else: out_data = self.ws.send(TextMessage(data=str(data))) self.sock.send(out_data) def receive(self, timeout=None): """Receive data over the WebSocket connection. :param timeout: Amount of time to wait for the data, in seconds. Set to ``None`` (the default) to wait indefinitely. Set to 0 to read without blocking. The data received is returned, as ``bytes`` or ``str``, depending on the type of the incoming message. """ while self.connected and not self.input_buffer: if not self.event.wait(timeout=timeout): return None self.event.clear() try: return self.input_buffer.pop(0) except IndexError: pass if not self.connected: # pragma: no cover raise ConnectionClosed(self.close_reason, self.close_message) def close(self, reason=None, message=None): """Close the WebSocket connection. :param reason: A numeric status code indicating the reason of the closure, as defined by the WebSocket specification. The default is 1000 (normal closure). :param message: A text message to be sent to the other side. """ if not self.connected: raise ConnectionClosed(self.close_reason, self.close_message) out_data = self.ws.send(CloseConnection( reason or CloseReason.NORMAL_CLOSURE, message)) try: self.sock.send(out_data) except BrokenPipeError: # pragma: no cover pass self.connected = False def choose_subprotocol(self, request): # pragma: no cover # The method should return the subprotocol to use, or ``None`` if no # subprotocol is chosen. Can be overridden by subclasses that implement # the server-side of the WebSocket protocol. return None def _thread(self): sel = None if self.ping_interval: next_ping = time() + self.ping_interval sel = self.selector_class() try: sel.register(self.sock, selectors.EVENT_READ, True) except ValueError: # pragma: no cover self.connected = False while self.connected: try: if sel: now = time() if next_ping <= now or not sel.select(next_ping - now): # we reached the timeout, we have to send a ping if not self.pong_received: self.close(reason=CloseReason.POLICY_VIOLATION, message='Ping/Pong timeout') self.event.set() break self.pong_received = False self.sock.send(self.ws.send(Ping())) next_ping = max(now, next_ping) + self.ping_interval continue in_data = self.sock.recv(self.receive_bytes) if len(in_data) == 0: raise OSError() self.ws.receive_data(in_data) self.connected = self._handle_events() except (OSError, ConnectionResetError, LocalProtocolError): # pragma: no cover self.connected = False self.event.set() break sel.close() if sel else None self.sock.close() def _handle_events(self): keep_going = True out_data = b'' for event in self.ws.events(): try: if isinstance(event, Request): self.subprotocol = self.choose_subprotocol(event) out_data += self.ws.send(AcceptConnection( subprotocol=self.subprotocol, extensions=[PerMessageDeflate()])) elif isinstance(event, CloseConnection): if self.is_server: out_data += self.ws.send(event.response()) self.close_reason = event.code self.close_message = event.reason self.connected = False self.event.set() keep_going = False elif isinstance(event, Ping): out_data += self.ws.send(event.response()) elif isinstance(event, Pong): self.pong_received = True elif isinstance(event, (TextMessage, BytesMessage)): self.incoming_message_len += len(event.data) if self.max_message_size and \ self.incoming_message_len > self.max_message_size: out_data += self.ws.send(CloseConnection( CloseReason.MESSAGE_TOO_BIG, 'Message is too big')) self.event.set() keep_going = False break if self.incoming_message is None: # store message as is first # if it is the first of a group, the message will be # converted to bytearray on arrival of the second # part, since bytearrays are mutable and can be # concatenated more efficiently self.incoming_message = event.data elif isinstance(event, TextMessage): if not isinstance(self.incoming_message, bytearray): # convert to bytearray and append self.incoming_message = bytearray( (self.incoming_message + event.data).encode()) else: # append to bytearray self.incoming_message += event.data.encode() else: if not isinstance(self.incoming_message, bytearray): # convert to mutable bytearray and append self.incoming_message = bytearray( self.incoming_message + event.data) else: # append to bytearray self.incoming_message += event.data if not event.message_finished: continue if isinstance(self.incoming_message, (str, bytes)): # single part message self.input_buffer.append(self.incoming_message) elif isinstance(event, TextMessage): # convert multi-part message back to text self.input_buffer.append( self.incoming_message.decode()) else: # convert multi-part message back to bytes self.input_buffer.append(bytes(self.incoming_message)) self.incoming_message = None self.incoming_message_len = 0 self.event.set() else: # pragma: no cover pass except LocalProtocolError: # pragma: no cover out_data = b'' self.event.set() keep_going = False if out_data: self.sock.send(out_data) return keep_going class Server(Base): """This class implements a WebSocket server. Instead of creating an instance of this class directly, use the ``accept()`` class method to create individual instances of the server, each bound to a client request. """ def __init__(self, environ, subprotocols=None, receive_bytes=4096, ping_interval=None, max_message_size=None, thread_class=None, event_class=None, selector_class=None): self.environ = environ self.subprotocols = subprotocols or [] if isinstance(self.subprotocols, str): self.subprotocols = [self.subprotocols] self.mode = 'unknown' sock = None if 'werkzeug.socket' in environ: # extract socket from Werkzeug's WSGI environment sock = environ.get('werkzeug.socket') self.mode = 'werkzeug' elif 'gunicorn.socket' in environ: # extract socket from Gunicorn WSGI environment sock = environ.get('gunicorn.socket') self.mode = 'gunicorn' elif 'eventlet.input' in environ: # pragma: no cover # extract socket from Eventlet's WSGI environment sock = environ.get('eventlet.input').get_socket() self.mode = 'eventlet' elif environ.get('SERVER_SOFTWARE', '').startswith( 'gevent'): # pragma: no cover # extract socket from Gevent's WSGI environment wsgi_input = environ['wsgi.input'] if not hasattr(wsgi_input, 'raw') and hasattr(wsgi_input, 'rfile'): wsgi_input = wsgi_input.rfile if hasattr(wsgi_input, 'raw'): sock = wsgi_input.raw._sock try: sock = sock.dup() except NotImplementedError: pass self.mode = 'gevent' if sock is None: raise RuntimeError('Cannot obtain socket from WSGI environment.') super().__init__(sock, connection_type=ConnectionType.SERVER, receive_bytes=receive_bytes, ping_interval=ping_interval, max_message_size=max_message_size, thread_class=thread_class, event_class=event_class, selector_class=selector_class) @classmethod def accept(cls, environ, subprotocols=None, receive_bytes=4096, ping_interval=None, max_message_size=None, thread_class=None, event_class=None, selector_class=None): """Accept a WebSocket connection from a client. :param environ: A WSGI ``environ`` dictionary with the request details. Among other things, this class expects to find the low-level network socket for the connection somewhere in this dictionary. Since the WSGI specification does not cover where or how to store this socket, each web server does this in its own different way. Werkzeug, Gunicorn, Eventlet and Gevent are the only web servers that are currently supported. :param subprotocols: A list of supported subprotocols, or ``None`` (the default) to disable subprotocol negotiation. :param receive_bytes: The size of the receive buffer, in bytes. The default is 4096. :param ping_interval: Send ping packets to clients at the requested interval in seconds. Set to ``None`` (the default) to disable ping/pong logic. Enable to prevent disconnections when the line is idle for a certain amount of time, or to detect unresponsive clients and disconnect them. A recommended interval is 25 seconds. :param max_message_size: The maximum size allowed for a message, in bytes, or ``None`` for no limit. The default is ``None``. :param thread_class: The ``Thread`` class to use when creating background threads. The default is the ``threading.Thread`` class from the Python standard library. :param event_class: The ``Event`` class to use when creating event objects. The default is the `threading.Event`` class from the Python standard library. :param selector_class: The ``Selector`` class to use when creating selectors. The default is the ``selectors.DefaultSelector`` class from the Python standard library. """ return cls(environ, subprotocols=subprotocols, receive_bytes=receive_bytes, ping_interval=ping_interval, max_message_size=max_message_size, thread_class=thread_class, event_class=event_class, selector_class=selector_class) def handshake(self): in_data = b'GET / HTTP/1.1\r\n' for key, value in self.environ.items(): if key.startswith('HTTP_'): header = '-'.join([p.capitalize() for p in key[5:].split('_')]) in_data += f'{header}: {value}\r\n'.encode() in_data += b'\r\n' self.ws.receive_data(in_data) self.connected = self._handle_events() def choose_subprotocol(self, request): """Choose a subprotocol to use for the WebSocket connection. The default implementation selects the first protocol requested by the client that is accepted by the server. Subclasses can override this method to implement a different subprotocol negotiation algorithm. :param request: A ``Request`` object. The method should return the subprotocol to use, or ``None`` if no subprotocol is chosen. """ for subprotocol in request.subprotocols: if subprotocol in self.subprotocols: return subprotocol return None class Client(Base): """This class implements a WebSocket client. Instead of creating an instance of this class directly, use the ``connect()`` class method to create an instance that is connected to a server. """ def __init__(self, url, subprotocols=None, headers=None, receive_bytes=4096, ping_interval=None, max_message_size=None, ssl_context=None, thread_class=None, event_class=None): parsed_url = urlsplit(url) is_secure = parsed_url.scheme in ['https', 'wss'] self.host = parsed_url.hostname self.port = parsed_url.port or (443 if is_secure else 80) self.path = parsed_url.path if parsed_url.query: self.path += '?' + parsed_url.query self.subprotocols = subprotocols or [] if isinstance(self.subprotocols, str): self.subprotocols = [self.subprotocols] self.extra_headeers = [] if isinstance(headers, dict): for key, value in headers.items(): self.extra_headeers.append((key, value)) elif isinstance(headers, list): self.extra_headeers = headers connection_args = socket.getaddrinfo(self.host, self.port, type=socket.SOCK_STREAM) if len(connection_args) == 0: # pragma: no cover raise ConnectionError() sock = socket.socket(connection_args[0][0], connection_args[0][1], connection_args[0][2]) if is_secure: # pragma: no cover if ssl_context is None: ssl_context = ssl.create_default_context( purpose=ssl.Purpose.SERVER_AUTH) sock = ssl_context.wrap_socket(sock, server_hostname=self.host) sock.connect(connection_args[0][4]) super().__init__(sock, connection_type=ConnectionType.CLIENT, receive_bytes=receive_bytes, ping_interval=ping_interval, max_message_size=max_message_size, thread_class=thread_class, event_class=event_class) @classmethod def connect(cls, url, subprotocols=None, headers=None, receive_bytes=4096, ping_interval=None, max_message_size=None, ssl_context=None, thread_class=None, event_class=None): """Returns a WebSocket client connection. :param url: The connection URL. Both ``ws://`` and ``wss://`` URLs are accepted. :param subprotocols: The name of the subprotocol to use, or a list of subprotocol names in order of preference. Set to ``None`` (the default) to not use a subprotocol. :param headers: A dictionary or list of tuples with additional HTTP headers to send with the connection request. Note that custom headers are not supported by the WebSocket protocol, so the use of this parameter is not recommended. :param receive_bytes: The size of the receive buffer, in bytes. The default is 4096. :param ping_interval: Send ping packets to the server at the requested interval in seconds. Set to ``None`` (the default) to disable ping/pong logic. Enable to prevent disconnections when the line is idle for a certain amount of time, or to detect an unresponsive server and disconnect. A recommended interval is 25 seconds. In general it is preferred to enable ping/pong on the server, and let the client respond with pong (which it does regardless of this setting). :param max_message_size: The maximum size allowed for a message, in bytes, or ``None`` for no limit. The default is ``None``. :param ssl_context: An ``SSLContext`` instance, if a default SSL context isn't sufficient. :param thread_class: The ``Thread`` class to use when creating background threads. The default is the ``threading.Thread`` class from the Python standard library. :param event_class: The ``Event`` class to use when creating event objects. The default is the `threading.Event`` class from the Python standard library. """ return cls(url, subprotocols=subprotocols, headers=headers, receive_bytes=receive_bytes, ping_interval=ping_interval, max_message_size=max_message_size, ssl_context=ssl_context, thread_class=thread_class, event_class=event_class) def handshake(self): out_data = self.ws.send(Request(host=self.host, target=self.path, subprotocols=self.subprotocols, extra_headers=self.extra_headeers)) self.sock.send(out_data) while True: in_data = self.sock.recv(self.receive_bytes) self.ws.receive_data(in_data) try: event = next(self.ws.events()) except StopIteration: # pragma: no cover pass else: # pragma: no cover break if isinstance(event, RejectConnection): # pragma: no cover raise ConnectionError(event.status_code) elif not isinstance(event, AcceptConnection): # pragma: no cover raise ConnectionError(400) self.subprotocol = event.subprotocol self.connected = True def close(self, reason=None, message=None): super().close(reason=reason, message=message) self.sock.close() simple-websocket-1.1.0/src/simple_websocket/aiows.py0000664000175000017500000005076214507561517017557 0ustar import asyncio import ssl from time import time from urllib.parse import urlsplit from wsproto import ConnectionType, WSConnection from wsproto.events import ( AcceptConnection, RejectConnection, CloseConnection, Message, Request, Ping, Pong, TextMessage, BytesMessage, ) from wsproto.extensions import PerMessageDeflate from wsproto.frame_protocol import CloseReason from wsproto.utilities import LocalProtocolError from .errors import ConnectionError, ConnectionClosed class AioBase: def __init__(self, connection_type=None, receive_bytes=4096, ping_interval=None, max_message_size=None): #: The name of the subprotocol chosen for the WebSocket connection. self.subprotocol = None self.connection_type = connection_type self.receive_bytes = receive_bytes self.ping_interval = ping_interval self.max_message_size = max_message_size self.pong_received = True self.input_buffer = [] self.incoming_message = None self.incoming_message_len = 0 self.connected = False self.is_server = (connection_type == ConnectionType.SERVER) self.close_reason = CloseReason.NO_STATUS_RCVD self.close_message = None self.rsock = None self.wsock = None self.event = asyncio.Event() self.ws = None self.task = None async def connect(self): self.ws = WSConnection(self.connection_type) await self.handshake() if not self.connected: # pragma: no cover raise ConnectionError() self.task = asyncio.create_task(self._task()) async def handshake(self): # pragma: no cover # to be implemented by subclasses pass async def send(self, data): """Send data over the WebSocket connection. :param data: The data to send. If ``data`` is of type ``bytes``, then a binary message is sent. Else, the message is sent in text format. """ if not self.connected: raise ConnectionClosed(self.close_reason, self.close_message) if isinstance(data, bytes): out_data = self.ws.send(Message(data=data)) else: out_data = self.ws.send(TextMessage(data=str(data))) self.wsock.write(out_data) async def receive(self, timeout=None): """Receive data over the WebSocket connection. :param timeout: Amount of time to wait for the data, in seconds. Set to ``None`` (the default) to wait indefinitely. Set to 0 to read without blocking. The data received is returned, as ``bytes`` or ``str``, depending on the type of the incoming message. """ while self.connected and not self.input_buffer: try: await asyncio.wait_for(self.event.wait(), timeout=timeout) except asyncio.TimeoutError: return None self.event.clear() # pragma: no cover try: return self.input_buffer.pop(0) except IndexError: pass if not self.connected: # pragma: no cover raise ConnectionClosed(self.close_reason, self.close_message) async def close(self, reason=None, message=None): """Close the WebSocket connection. :param reason: A numeric status code indicating the reason of the closure, as defined by the WebSocket specification. The default is 1000 (normal closure). :param message: A text message to be sent to the other side. """ if not self.connected: raise ConnectionClosed(self.close_reason, self.close_message) out_data = self.ws.send(CloseConnection( reason or CloseReason.NORMAL_CLOSURE, message)) try: self.wsock.write(out_data) except BrokenPipeError: # pragma: no cover pass self.connected = False def choose_subprotocol(self, request): # pragma: no cover # The method should return the subprotocol to use, or ``None`` if no # subprotocol is chosen. Can be overridden by subclasses that implement # the server-side of the WebSocket protocol. return None async def _task(self): next_ping = None if self.ping_interval: next_ping = time() + self.ping_interval while self.connected: try: in_data = b'' if next_ping: now = time() timed_out = True if next_ping > now: timed_out = False try: in_data = await asyncio.wait_for( self.rsock.read(self.receive_bytes), timeout=next_ping - now) except asyncio.TimeoutError: timed_out = True if timed_out: # we reached the timeout, we have to send a ping if not self.pong_received: await self.close( reason=CloseReason.POLICY_VIOLATION, message='Ping/Pong timeout') break self.pong_received = False self.wsock.write(self.ws.send(Ping())) next_ping = max(now, next_ping) + self.ping_interval continue else: in_data = await self.rsock.read(self.receive_bytes) if len(in_data) == 0: raise OSError() except (OSError, ConnectionResetError): # pragma: no cover self.connected = False self.event.set() break self.ws.receive_data(in_data) self.connected = await self._handle_events() self.wsock.close() async def _handle_events(self): keep_going = True out_data = b'' for event in self.ws.events(): try: if isinstance(event, Request): self.subprotocol = self.choose_subprotocol(event) out_data += self.ws.send(AcceptConnection( subprotocol=self.subprotocol, extensions=[PerMessageDeflate()])) elif isinstance(event, CloseConnection): if self.is_server: out_data += self.ws.send(event.response()) self.close_reason = event.code self.close_message = event.reason self.connected = False self.event.set() keep_going = False elif isinstance(event, Ping): out_data += self.ws.send(event.response()) elif isinstance(event, Pong): self.pong_received = True elif isinstance(event, (TextMessage, BytesMessage)): self.incoming_message_len += len(event.data) if self.max_message_size and \ self.incoming_message_len > self.max_message_size: out_data += self.ws.send(CloseConnection( CloseReason.MESSAGE_TOO_BIG, 'Message is too big')) self.event.set() keep_going = False break if self.incoming_message is None: # store message as is first # if it is the first of a group, the message will be # converted to bytearray on arrival of the second # part, since bytearrays are mutable and can be # concatenated more efficiently self.incoming_message = event.data elif isinstance(event, TextMessage): if not isinstance(self.incoming_message, bytearray): # convert to bytearray and append self.incoming_message = bytearray( (self.incoming_message + event.data).encode()) else: # append to bytearray self.incoming_message += event.data.encode() else: if not isinstance(self.incoming_message, bytearray): # convert to mutable bytearray and append self.incoming_message = bytearray( self.incoming_message + event.data) else: # append to bytearray self.incoming_message += event.data if not event.message_finished: continue if isinstance(self.incoming_message, (str, bytes)): # single part message self.input_buffer.append(self.incoming_message) elif isinstance(event, TextMessage): # convert multi-part message back to text self.input_buffer.append( self.incoming_message.decode()) else: # convert multi-part message back to bytes self.input_buffer.append(bytes(self.incoming_message)) self.incoming_message = None self.incoming_message_len = 0 self.event.set() else: # pragma: no cover pass except LocalProtocolError: # pragma: no cover out_data = b'' self.event.set() keep_going = False if out_data: self.wsock.write(out_data) return keep_going class AioServer(AioBase): """This class implements a WebSocket server. Instead of creating an instance of this class directly, use the ``accept()`` class method to create individual instances of the server, each bound to a client request. """ def __init__(self, request, subprotocols=None, receive_bytes=4096, ping_interval=None, max_message_size=None): super().__init__(connection_type=ConnectionType.SERVER, receive_bytes=receive_bytes, ping_interval=ping_interval, max_message_size=max_message_size) self.request = request self.headers = {} self.subprotocols = subprotocols or [] if isinstance(self.subprotocols, str): self.subprotocols = [self.subprotocols] self.mode = 'unknown' @classmethod async def accept(cls, aiohttp=None, asgi=None, sock=None, headers=None, subprotocols=None, receive_bytes=4096, ping_interval=None, max_message_size=None): """Accept a WebSocket connection from a client. :param aiohttp: The request object from aiohttp. If this argument is provided, ``asgi``, ``sock`` and ``headers`` must not be set. :param asgi: A (scope, receive, send) tuple from an ASGI request. If this argument is provided, ``aiohttp``, ``sock`` and ``headers`` must not be set. :param sock: A connected socket to use. If this argument is provided, ``aiohttp`` and ``asgi`` must not be set. The ``headers`` argument must be set with the incoming request headers. :param headers: A dictionary with the incoming request headers, when ``sock`` is used. :param subprotocols: A list of supported subprotocols, or ``None`` (the default) to disable subprotocol negotiation. :param receive_bytes: The size of the receive buffer, in bytes. The default is 4096. :param ping_interval: Send ping packets to clients at the requested interval in seconds. Set to ``None`` (the default) to disable ping/pong logic. Enable to prevent disconnections when the line is idle for a certain amount of time, or to detect unresponsive clients and disconnect them. A recommended interval is 25 seconds. :param max_message_size: The maximum size allowed for a message, in bytes, or ``None`` for no limit. The default is ``None``. """ if aiohttp and (asgi or sock): raise ValueError('aiohttp argument cannot be used with asgi or ' 'sock') if asgi and (aiohttp or sock): raise ValueError('asgi argument cannot be used with aiohttp or ' 'sock') if asgi: # pragma: no cover from .asgi import WebSocketASGI return await WebSocketASGI.accept(asgi[0], asgi[1], asgi[2], subprotocols=subprotocols) ws = cls({'aiohttp': aiohttp, 'sock': sock, 'headers': headers}, subprotocols=subprotocols, receive_bytes=receive_bytes, ping_interval=ping_interval, max_message_size=max_message_size) await ws._accept() return ws async def _accept(self): if self.request['sock']: # pragma: no cover # custom integration, request is a tuple with (socket, headers) sock = self.request['sock'] self.headers = self.request['headers'] self.mode = 'custom' elif self.request['aiohttp']: # default implementation, request is an aiohttp request object sock = self.request['aiohttp'].transport.get_extra_info( 'socket').dup() self.headers = self.request['aiohttp'].headers self.mode = 'aiohttp' else: # pragma: no cover raise ValueError('Invalid request') self.rsock, self.wsock = await asyncio.open_connection(sock=sock) await super().connect() async def handshake(self): in_data = b'GET / HTTP/1.1\r\n' for header, value in self.headers.items(): in_data += f'{header}: {value}\r\n'.encode() in_data += b'\r\n' self.ws.receive_data(in_data) self.connected = await self._handle_events() def choose_subprotocol(self, request): """Choose a subprotocol to use for the WebSocket connection. The default implementation selects the first protocol requested by the client that is accepted by the server. Subclasses can override this method to implement a different subprotocol negotiation algorithm. :param request: A ``Request`` object. The method should return the subprotocol to use, or ``None`` if no subprotocol is chosen. """ for subprotocol in request.subprotocols: if subprotocol in self.subprotocols: return subprotocol return None class AioClient(AioBase): """This class implements a WebSocket client. Instead of creating an instance of this class directly, use the ``connect()`` class method to create an instance that is connected to a server. """ def __init__(self, url, subprotocols=None, headers=None, receive_bytes=4096, ping_interval=None, max_message_size=None, ssl_context=None): super().__init__(connection_type=ConnectionType.CLIENT, receive_bytes=receive_bytes, ping_interval=ping_interval, max_message_size=max_message_size) self.url = url self.ssl_context = ssl_context parsed_url = urlsplit(url) self.is_secure = parsed_url.scheme in ['https', 'wss'] self.host = parsed_url.hostname self.port = parsed_url.port or (443 if self.is_secure else 80) self.path = parsed_url.path if parsed_url.query: self.path += '?' + parsed_url.query self.subprotocols = subprotocols or [] if isinstance(self.subprotocols, str): self.subprotocols = [self.subprotocols] self.extra_headeers = [] if isinstance(headers, dict): for key, value in headers.items(): self.extra_headeers.append((key, value)) elif isinstance(headers, list): self.extra_headeers = headers @classmethod async def connect(cls, url, subprotocols=None, headers=None, receive_bytes=4096, ping_interval=None, max_message_size=None, ssl_context=None, thread_class=None, event_class=None): """Returns a WebSocket client connection. :param url: The connection URL. Both ``ws://`` and ``wss://`` URLs are accepted. :param subprotocols: The name of the subprotocol to use, or a list of subprotocol names in order of preference. Set to ``None`` (the default) to not use a subprotocol. :param headers: A dictionary or list of tuples with additional HTTP headers to send with the connection request. Note that custom headers are not supported by the WebSocket protocol, so the use of this parameter is not recommended. :param receive_bytes: The size of the receive buffer, in bytes. The default is 4096. :param ping_interval: Send ping packets to the server at the requested interval in seconds. Set to ``None`` (the default) to disable ping/pong logic. Enable to prevent disconnections when the line is idle for a certain amount of time, or to detect an unresponsive server and disconnect. A recommended interval is 25 seconds. In general it is preferred to enable ping/pong on the server, and let the client respond with pong (which it does regardless of this setting). :param max_message_size: The maximum size allowed for a message, in bytes, or ``None`` for no limit. The default is ``None``. :param ssl_context: An ``SSLContext`` instance, if a default SSL context isn't sufficient. """ ws = cls(url, subprotocols=subprotocols, headers=headers, receive_bytes=receive_bytes, ping_interval=ping_interval, max_message_size=max_message_size, ssl_context=ssl_context) await ws._connect() return ws async def _connect(self): if self.is_secure: # pragma: no cover if self.ssl_context is None: self.ssl_context = ssl.create_default_context( purpose=ssl.Purpose.SERVER_AUTH) self.rsock, self.wsock = await asyncio.open_connection( self.host, self.port, ssl=self.ssl_context) await super().connect() async def handshake(self): out_data = self.ws.send(Request(host=self.host, target=self.path, subprotocols=self.subprotocols, extra_headers=self.extra_headeers)) self.wsock.write(out_data) while True: in_data = await self.rsock.read(self.receive_bytes) self.ws.receive_data(in_data) try: event = next(self.ws.events()) except StopIteration: # pragma: no cover pass else: # pragma: no cover break if isinstance(event, RejectConnection): # pragma: no cover raise ConnectionError(event.status_code) elif not isinstance(event, AcceptConnection): # pragma: no cover raise ConnectionError(400) self.subprotocol = event.subprotocol self.connected = True async def close(self, reason=None, message=None): await super().close(reason=reason, message=message) self.wsock.close() simple-websocket-1.1.0/src/simple_websocket/__init__.py0000664000175000017500000000024714507561517020165 0ustar from .ws import Server, Client # noqa: F401 from .aiows import AioServer, AioClient # noqa: F401 from .errors import ConnectionError, ConnectionClosed # noqa: F401 simple-websocket-1.1.0/src/simple_websocket/errors.py0000664000175000017500000000121414702046370017725 0ustar from wsproto.frame_protocol import CloseReason class SimpleWebsocketError(RuntimeError): pass class ConnectionError(SimpleWebsocketError): """Connection error exception class.""" def __init__(self, status_code=None): # pragma: no cover self.status_code = status_code super().__init__(f'Connection error: {status_code}') class ConnectionClosed(SimpleWebsocketError): """Connection closed exception class.""" def __init__(self, reason=CloseReason.NO_STATUS_RCVD, message=None): self.reason = reason self.message = message super().__init__(f'Connection closed: {reason} {message or ""}') simple-websocket-1.1.0/src/simple_websocket/asgi.py0000664000175000017500000000343714507561517017355 0ustar from .errors import ConnectionClosed # pragma: no cover class WebSocketASGI: # pragma: no cover def __init__(self, scope, receive, send, subprotocols=None): self._scope = scope self._receive = receive self._send = send self.subprotocols = subprotocols or [] self.subprotocol = None self.connected = False @classmethod async def accept(cls, scope, receive, send, subprotocols=None): ws = WebSocketASGI(scope, receive, send, subprotocols=subprotocols) await ws._accept() return ws async def _accept(self): connect = await self._receive() if connect['type'] != 'websocket.connect': raise ValueError('Expected websocket.connect') for subprotocol in self._scope['subprotocols']: if subprotocol in self.subprotocols: self.subprotocol = subprotocol break await self._send({'type': 'websocket.accept', 'subprotocol': self.subprotocol}) async def receive(self): message = await self._receive() if message['type'] == 'websocket.disconnect': raise ConnectionClosed() elif message['type'] != 'websocket.receive': raise OSError(32, 'Websocket message type not supported') return message.get('text', message.get('bytes')) async def send(self, data): if isinstance(data, str): await self._send({'type': 'websocket.send', 'text': data}) else: await self._send({'type': 'websocket.send', 'bytes': data}) async def close(self): if not self.connected: self.conncted = False try: await self._send({'type': 'websocket.close'}) except Exception: pass simple-websocket-1.1.0/src/simple_websocket.egg-info/0000775000175000017500000000000014702053632017532 5ustar simple-websocket-1.1.0/src/simple_websocket.egg-info/not-zip-safe0000664000175000017500000000000114510045601021752 0ustar simple-websocket-1.1.0/src/simple_websocket.egg-info/top_level.txt0000664000175000017500000000002114702053632022255 0ustar simple_websocket simple-websocket-1.1.0/src/simple_websocket.egg-info/SOURCES.txt0000664000175000017500000000127514702053632021423 0ustar LICENSE MANIFEST.in README.md pyproject.toml tox.ini docs/Makefile docs/api.rst docs/conf.py docs/index.rst docs/intro.rst docs/make.bat docs/_static/css/custom.css src/simple_websocket/__init__.py src/simple_websocket/aiows.py src/simple_websocket/asgi.py src/simple_websocket/errors.py src/simple_websocket/ws.py src/simple_websocket.egg-info/PKG-INFO src/simple_websocket.egg-info/SOURCES.txt src/simple_websocket.egg-info/dependency_links.txt src/simple_websocket.egg-info/not-zip-safe src/simple_websocket.egg-info/requires.txt src/simple_websocket.egg-info/top_level.txt tests/__init__.py tests/helpers.py tests/test_aioclient.py tests/test_aioserver.py tests/test_client.py tests/test_server.pysimple-websocket-1.1.0/src/simple_websocket.egg-info/dependency_links.txt0000664000175000017500000000000114702053632023600 0ustar simple-websocket-1.1.0/src/simple_websocket.egg-info/requires.txt0000664000175000017500000000007314702053632022132 0ustar wsproto [dev] tox flake8 pytest pytest-cov [docs] sphinx simple-websocket-1.1.0/src/simple_websocket.egg-info/PKG-INFO0000644000175000017500000000276514702053632020637 0ustar Metadata-Version: 2.1 Name: simple-websocket Version: 1.1.0 Summary: Simple WebSocket server and client for Python Author-email: Miguel Grinberg Project-URL: Homepage, https://github.com/miguelgrinberg/simple-websocket Project-URL: Bug Tracker, https://github.com/miguelgrinberg/simple-websocket/issues Classifier: Environment :: Web Environment Classifier: Intended Audience :: Developers Classifier: Programming Language :: Python :: 3 Classifier: License :: OSI Approved :: MIT License Classifier: Operating System :: OS Independent Requires-Python: >=3.6 Description-Content-Type: text/markdown License-File: LICENSE Requires-Dist: wsproto Provides-Extra: dev Requires-Dist: tox; extra == "dev" Requires-Dist: flake8; extra == "dev" Requires-Dist: pytest; extra == "dev" Requires-Dist: pytest-cov; extra == "dev" Provides-Extra: docs Requires-Dist: sphinx; extra == "docs" simple-websocket ================ [![Build status](https://github.com/miguelgrinberg/simple-websocket/workflows/build/badge.svg)](https://github.com/miguelgrinberg/simple-websocket/actions) [![codecov](https://codecov.io/gh/miguelgrinberg/simple-websocket/branch/main/graph/badge.svg)](https://codecov.io/gh/miguelgrinberg/simple-websocket) Simple WebSocket server and client for Python. ## Resources - [Documentation](http://simple-websocket.readthedocs.io/en/latest/) - [PyPI](https://pypi.python.org/pypi/simple-websocket) - [Change Log](https://github.com/miguelgrinberg/simple-websocket/blob/main/CHANGES.md) simple-websocket-1.1.0/README.md0000664000175000017500000000115314444310476013176 0ustar simple-websocket ================ [![Build status](https://github.com/miguelgrinberg/simple-websocket/workflows/build/badge.svg)](https://github.com/miguelgrinberg/simple-websocket/actions) [![codecov](https://codecov.io/gh/miguelgrinberg/simple-websocket/branch/main/graph/badge.svg)](https://codecov.io/gh/miguelgrinberg/simple-websocket) Simple WebSocket server and client for Python. ## Resources - [Documentation](http://simple-websocket.readthedocs.io/en/latest/) - [PyPI](https://pypi.python.org/pypi/simple-websocket) - [Change Log](https://github.com/miguelgrinberg/simple-websocket/blob/main/CHANGES.md) simple-websocket-1.1.0/docs/0000775000175000017500000000000014702053632012642 5ustar simple-websocket-1.1.0/docs/intro.rst0000664000175000017500000000716414507561517014550 0ustar Getting Started =============== ``simple-websocket`` includes a collection of WebSocket servers and clients for Python, including support for both traditional and asynchronous (asyncio) workflows. The servers are designed to be integrated into larger web applications if desired. Installation ------------ This package is installed with ``pip``:: pip install simple-websocket Server Example #1: Flask ------------------------ The following example shows how to add a WebSocket route to a `Flask `_ application. :: from flask import Flask, request from simple_websocket import Server, ConnectionClosed app = Flask(__name__) @app.route('/echo', websocket=True) def echo(): ws = Server.accept(request.environ) try: while True: data = ws.receive() ws.send(data) except ConnectionClosed: pass return '' Integration with web applications using other `WSGI `_ frameworks works in a similar way. The only requirement is to pass the ``environ`` dictionary to the ``Server.accept()`` method to initiate the WebSocket handshake. Server Example #2: Aiohttp -------------------------- The following example shows how to add a WebSocket route to a web application built with the `aiohttp `_ framework. :: from aiohttp import web from simple_websocket import AioServer, ConnectionClosed app = web.Application() async def echo(request): ws = await AioServer.accept(aiohttp=request) try: while True: data = await ws.receive() await ws.send(data) except ConnectionClosed: pass return web.Response(text='') app.add_routes([web.get('/echo', echo)]) if __name__ == '__main__': web.run_app(app, port=5000) Server Example #3: ASGI ----------------------- The next server example shows an asynchronous application that supports the `ASGI `_ protocol. :: from simple_websocket import AioServer, ConnectionClosed async def echo(scope, receive, send): ws = await AioServer.accept(asgi=(scope, receive, send)) try: while True: data = await ws.receive() await ws.send(data) except ConnectionClosed: pass Client Example #1: Synchronous ------------------------------ The client example that follows can connect to any of the server examples above using a synchronous interface. :: from simple_websocket import Client, ConnectionClosed def main(): ws = Client.connect('ws://localhost:5000/echo') try: while True: data = input('> ') ws.send(data) data = ws.receive() print(f'< {data}') except (KeyboardInterrupt, EOFError, ConnectionClosed): ws.close() if __name__ == '__main__': main() Client Example #2: Asynchronous ------------------------------- The next client uses Python's ``asyncio`` framework. :: import asyncio from simple_websocket import AioClient, ConnectionClosed async def main(): ws = await AioClient.connect('ws://localhost:5000/echo') try: while True: data = input('> ') await ws.send(data) data = await ws.receive() print(f'< {data}') except (KeyboardInterrupt, EOFError, ConnectionClosed): await ws.close() if __name__ == '__main__': asyncio.run(main()) simple-websocket-1.1.0/docs/index.rst0000664000175000017500000000056514444310476014516 0ustar .. simple-websocket documentation master file, created by sphinx-quickstart on Mon Jun 7 14:14:04 2021. You can adapt this file completely to your liking, but it should at least contain the root `toctree` directive. simple-websocket ================ Simple WebSocket server and client for Python. .. toctree:: :maxdepth: 2 intro api * :ref:`search` simple-websocket-1.1.0/docs/make.bat0000664000175000017500000000143314444310476014255 0ustar @ECHO OFF pushd %~dp0 REM Command file for Sphinx documentation if "%SPHINXBUILD%" == "" ( set SPHINXBUILD=sphinx-build ) set SOURCEDIR=. set BUILDDIR=_build if "%1" == "" goto help %SPHINXBUILD% >NUL 2>NUL if errorlevel 9009 ( echo. echo.The 'sphinx-build' command was not found. Make sure you have Sphinx echo.installed, then set the SPHINXBUILD environment variable to point echo.to the full path of the 'sphinx-build' executable. Alternatively you echo.may add the Sphinx directory to PATH. echo. echo.If you don't have Sphinx installed, grab it from echo.http://sphinx-doc.org/ exit /b 1 ) %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% goto end :help %SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% :end popd simple-websocket-1.1.0/docs/Makefile0000664000175000017500000000117214444310476014310 0ustar # Minimal makefile for Sphinx documentation # # You can set these variables from the command line, and also # from the environment for the first two. SPHINXOPTS ?= SPHINXBUILD ?= sphinx-build SOURCEDIR = . BUILDDIR = _build # Put it first so that "make" without argument is like "make help". help: @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) .PHONY: help Makefile # Catch-all target: route all unknown targets to Sphinx using the new # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). %: Makefile @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) simple-websocket-1.1.0/docs/api.rst0000664000175000017500000000123614507561517014160 0ustar API Reference ------------- The ``Server`` class ~~~~~~~~~~~~~~~~~~~~ .. autoclass:: simple_websocket.Server :inherited-members: :members: The ``AioServer`` class ~~~~~~~~~~~~~~~~~~~~~~~ .. autoclass:: simple_websocket.AioServer :inherited-members: :members: The ``Client`` class ~~~~~~~~~~~~~~~~~~~~ .. autoclass:: simple_websocket.Client :inherited-members: :members: The ``AioClient`` class ~~~~~~~~~~~~~~~~~~~~~~~ .. autoclass:: simple_websocket.AioClient :inherited-members: :members: Exceptions ~~~~~~~~~~ .. autoclass:: simple_websocket.ConnectionError :members: .. autoclass:: simple_websocket.ConnectionClosed :members: simple-websocket-1.1.0/docs/conf.py0000664000175000017500000000421114507561517014150 0ustar # Configuration file for the Sphinx documentation builder. # # This file only contains a selection of the most common options. For a full # list see the documentation: # https://www.sphinx-doc.org/en/master/usage/configuration.html # -- Path setup -------------------------------------------------------------- # If extensions (or modules to document with autodoc) are in another directory, # add these directories to sys.path here. If the directory is relative to the # documentation root, use os.path.abspath to make it absolute, like shown here. # import os import sys sys.path.insert(0, os.path.abspath('../src')) # -- Project information ----------------------------------------------------- project = 'simple-websocket' copyright = '2021, Miguel Grinberg' author = 'Miguel Grinberg' # -- General configuration --------------------------------------------------- # Add any Sphinx extension module names here, as strings. They can be # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom # ones. extensions = [ 'sphinx.ext.autodoc', ] autodoc_member_order = 'bysource' # Add any paths that contain templates here, relative to this directory. templates_path = ['_templates'] # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. # This pattern also affects html_static_path and html_extra_path. exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] # -- Options for HTML output ------------------------------------------------- # The theme to use for HTML and HTML Help pages. See the documentation for # a list of builtin themes. # html_theme = 'alabaster' # Add any paths that contain custom static files (such as style sheets) here, # relative to this directory. They are copied after the builtin static files, # so a file named "default.css" will overwrite the builtin "default.css". html_static_path = ['_static'] html_css_files = [ 'css/custom.css', ] html_theme_options = { 'github_user': 'miguelgrinberg', 'github_repo': 'simple-websocket', 'github_banner': True, 'github_button': True, 'github_type': 'star', 'fixed_sidebar': True, } simple-websocket-1.1.0/docs/_static/0000775000175000017500000000000014702053632014270 5ustar simple-websocket-1.1.0/docs/_static/css/0000775000175000017500000000000014702053632015060 5ustar simple-websocket-1.1.0/docs/_static/css/custom.css0000664000175000017500000000007714444310476017115 0ustar .py .class, .py .method, .py .property { margin-top: 20px; } simple-websocket-1.1.0/setup.cfg0000664000175000017500000000004614702053632013533 0ustar [egg_info] tag_build = tag_date = 0 simple-websocket-1.1.0/LICENSE0000664000175000017500000000206014444310476012722 0ustar MIT License Copyright (c) 2021 Miguel Grinberg Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. simple-websocket-1.1.0/pyproject.toml0000664000175000017500000000213114702053630014621 0ustar [project] name = "simple-websocket" version = "1.1.0" authors = [ { name = "Miguel Grinberg", email = "miguel.grinberg@gmail.com" }, ] description = "Simple WebSocket server and client for Python" classifiers = [ "Environment :: Web Environment", "Intended Audience :: Developers", "Programming Language :: Python :: 3", "License :: OSI Approved :: MIT License", "Operating System :: OS Independent", ] requires-python = ">=3.6" dependencies = [ "wsproto", ] [project.readme] file = "README.md" content-type = "text/markdown" [project.urls] Homepage = "https://github.com/miguelgrinberg/simple-websocket" "Bug Tracker" = "https://github.com/miguelgrinberg/simple-websocket/issues" [project.optional-dependencies] dev = [ "tox", "flake8", "pytest", "pytest-cov", ] docs = [ "sphinx", ] [tool.setuptools] zip-safe = false include-package-data = true [tool.setuptools.package-dir] "" = "src" [tool.setuptools.packages.find] where = [ "src", ] namespaces = false [build-system] requires = [ "setuptools>=61.2", ] build-backend = "setuptools.build_meta" simple-websocket-1.1.0/PKG-INFO0000644000175000017500000000276514702053632013017 0ustar Metadata-Version: 2.1 Name: simple-websocket Version: 1.1.0 Summary: Simple WebSocket server and client for Python Author-email: Miguel Grinberg Project-URL: Homepage, https://github.com/miguelgrinberg/simple-websocket Project-URL: Bug Tracker, https://github.com/miguelgrinberg/simple-websocket/issues Classifier: Environment :: Web Environment Classifier: Intended Audience :: Developers Classifier: Programming Language :: Python :: 3 Classifier: License :: OSI Approved :: MIT License Classifier: Operating System :: OS Independent Requires-Python: >=3.6 Description-Content-Type: text/markdown License-File: LICENSE Requires-Dist: wsproto Provides-Extra: dev Requires-Dist: tox; extra == "dev" Requires-Dist: flake8; extra == "dev" Requires-Dist: pytest; extra == "dev" Requires-Dist: pytest-cov; extra == "dev" Provides-Extra: docs Requires-Dist: sphinx; extra == "docs" simple-websocket ================ [![Build status](https://github.com/miguelgrinberg/simple-websocket/workflows/build/badge.svg)](https://github.com/miguelgrinberg/simple-websocket/actions) [![codecov](https://codecov.io/gh/miguelgrinberg/simple-websocket/branch/main/graph/badge.svg)](https://codecov.io/gh/miguelgrinberg/simple-websocket) Simple WebSocket server and client for Python. ## Resources - [Documentation](http://simple-websocket.readthedocs.io/en/latest/) - [PyPI](https://pypi.python.org/pypi/simple-websocket) - [Change Log](https://github.com/miguelgrinberg/simple-websocket/blob/main/CHANGES.md) simple-websocket-1.1.0/tox.ini0000664000175000017500000000116114702053147013225 0ustar [tox] envlist=flake8,py38,py39,py310,py311,py312,py313,pypy3,docs skip_missing_interpreters=True [gh-actions] python = 3.8: py38 3.9: py39 3.10: py310 3.11: py311 3.12: py312 3.13: py313 pypy-3: pypy3 [testenv] commands= pip install -e . pytest -p no:logging --cov=simple_websocket --cov-branch --cov-report=term-missing --cov-report=xml deps= pytest pytest-cov [testenv:pypy3] [testenv:flake8] deps= flake8 commands= flake8 --exclude=".*" src/simple_websocket tests [testenv:docs] changedir=docs deps= sphinx allowlist_externals= make commands= make html simple-websocket-1.1.0/tests/0000775000175000017500000000000014702053632013054 5ustar simple-websocket-1.1.0/tests/test_aioserver.py0000664000175000017500000003140714507561517016502 0ustar import asyncio import unittest from unittest import mock import pytest # noqa: F401 from wsproto.events import Request, CloseConnection, TextMessage, \ BytesMessage, Ping, Pong import simple_websocket from .helpers import make_sync, AsyncMock class AioSimpleWebSocketServerTestCase(unittest.TestCase): async def get_server(self, mock_wsconn, request, events=[], client_subprotocols=None, server_subprotocols=None, **kwargs): mock_wsconn().events.side_effect = \ [iter(ev) for ev in [[ Request(host='example.com', target='/ws', subprotocols=client_subprotocols or [])]] + events + [[CloseConnection(1000, 'bye')]]] mock_wsconn().send = lambda x: str(x).encode('utf-8') request.headers.update({ 'Host': 'example.com', 'Connection': 'Upgrade', 'Upgrade': 'websocket', 'Sec-Websocket-Key': 'Iv8io/9s+lYFgZWcXczP8Q==', 'Sec-Websocket-Version': '13', }) return await simple_websocket.AioServer.accept( aiohttp=request, subprotocols=server_subprotocols, **kwargs) @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_aiohttp(self, mock_wsconn, mock_open_connection): mock_request = mock.MagicMock(headers={}) rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) server = await self.get_server(mock_wsconn, mock_request) assert server.rsock == rsock assert server.wsock == wsock assert server.mode == 'aiohttp' assert server.receive_bytes == 4096 assert server.input_buffer == [] assert server.event.__class__.__name__ == 'Event' mock_wsconn().receive_data.assert_any_call( b'GET / HTTP/1.1\r\n' b'Host: example.com\r\n' b'Connection: Upgrade\r\n' b'Upgrade: websocket\r\n' b'Sec-Websocket-Key: Iv8io/9s+lYFgZWcXczP8Q==\r\n' b'Sec-Websocket-Version: 13\r\n\r\n') assert server.is_server @make_sync async def test_invalid_request(self): with pytest.raises(ValueError): await simple_websocket.AioServer.accept(aiohttp='foo', asgi='bar') with pytest.raises(ValueError): await simple_websocket.AioServer.accept(asgi='bar', sock='baz') @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_send(self, mock_wsconn, mock_open_connection): mock_request = mock.MagicMock(headers={}) rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) server = await self.get_server(mock_wsconn, mock_request) while server.connected: await asyncio.sleep(0.01) with pytest.raises(simple_websocket.ConnectionClosed): await server.send('hello') server.connected = True await server.send('hello') wsock.write.assert_called_with( b"TextMessage(data='hello', frame_finished=True, " b"message_finished=True)") server.connected = True await server.send(b'hello') wsock.write.assert_called_with( b"Message(data=b'hello', frame_finished=True, " b"message_finished=True)") @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_receive(self, mock_wsconn, mock_open_connection): mock_request = mock.MagicMock(headers={}) rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) server = await self.get_server(mock_wsconn, mock_request, events=[ [TextMessage('hello')], [BytesMessage(b'hello')], ]) while server.connected: await asyncio.sleep(0.01) server.connected = True assert await server.receive() == 'hello' assert await server.receive() == b'hello' assert await server.receive(timeout=0) is None @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_receive_after_close(self, mock_wsconn, mock_open_connection): mock_request = mock.MagicMock(headers={}) rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) server = await self.get_server(mock_wsconn, mock_request, events=[ [TextMessage('hello')], ]) while server.connected: await asyncio.sleep(0.01) assert await server.receive() == 'hello' with pytest.raises(simple_websocket.ConnectionClosed): await server.receive() @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_receive_split_messages(self, mock_wsconn, mock_open_connection): mock_request = mock.MagicMock(headers={}) rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) server = await self.get_server(mock_wsconn, mock_request, events=[ [TextMessage('hel', message_finished=False)], [TextMessage('lo')], [TextMessage('he', message_finished=False)], [TextMessage('l', message_finished=False)], [TextMessage('lo')], [BytesMessage(b'hel', message_finished=False)], [BytesMessage(b'lo')], [BytesMessage(b'he', message_finished=False)], [BytesMessage(b'l', message_finished=False)], [BytesMessage(b'lo')], ]) while server.connected: await asyncio.sleep(0.01) server.connected = True assert await server.receive() == 'hello' assert await server.receive() == 'hello' assert await server.receive() == b'hello' assert await server.receive() == b'hello' assert await server.receive(timeout=0) is None @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_receive_ping(self, mock_wsconn, mock_open_connection): mock_request = mock.MagicMock(headers={}) rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) server = await self.get_server(mock_wsconn, mock_request, events=[ [Ping(b'hello')], ]) while server.connected: await asyncio.sleep(0.01) wsock.write.assert_any_call(b"Pong(payload=b'hello')") @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_receive_empty(self, mock_wsconn, mock_open_connection): mock_request = mock.MagicMock(headers={}) rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) server = await self.get_server(mock_wsconn, mock_request, events=[ [TextMessage('hello')], ]) while server.connected: await asyncio.sleep(0.01) server.connected = True assert await server.receive() == 'hello' assert await server.receive(timeout=0) is None @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_receive_large(self, mock_wsconn, mock_open_connection): mock_request = mock.MagicMock(headers={}) rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) server = await self.get_server(mock_wsconn, mock_request, events=[ [TextMessage('hello')], [TextMessage('hello1')], ], max_message_size=5) while server.connected: await asyncio.sleep(0.01) server.connected = True assert await server.receive() == 'hello' assert await server.receive(timeout=0) is None @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_close(self, mock_wsconn, mock_open_connection): mock_request = mock.MagicMock(headers={}) rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) server = await self.get_server(mock_wsconn, mock_request) while server.connected: await asyncio.sleep(0.01) with pytest.raises(simple_websocket.ConnectionClosed) as exc: await server.close() assert str(exc.value) == 'Connection closed: 1000 bye' server.connected = True await server.close() assert not server.connected wsock.write.assert_called_with( b'CloseConnection(code=, ' b'reason=None)') @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') @mock.patch('simple_websocket.aiows.time') @mock.patch('simple_websocket.aiows.asyncio.wait_for') async def test_ping_pong(self, mock_wait_for, mock_time, mock_wsconn, mock_open_connection): mock_request = mock.MagicMock(headers={}) rsock = mock.MagicMock(read=AsyncMock()) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) server = await self.get_server(mock_wsconn, mock_request, events=[ [TextMessage('hello')], [Pong()], ], ping_interval=25) mock_wait_for.side_effect = [b'x', b'x', asyncio.TimeoutError, asyncio.TimeoutError] mock_time.side_effect = [0, 1, 25.01, 25.02, 28, 52, 76] await server._task() assert wsock.write.call_count == 4 assert wsock.write.call_args_list[1][0][0].startswith(b'Ping') assert wsock.write.call_args_list[2][0][0].startswith(b'Ping') assert wsock.write.call_args_list[3][0][0].startswith(b'Close') @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_subprotocols(self, mock_wsconn, mock_open_connection): mock_request = mock.MagicMock(headers={}) rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) server = await self.get_server(mock_wsconn, mock_request, client_subprotocols=['foo', 'bar'], server_subprotocols='bar') while server.connected: await asyncio.sleep(0.01) assert server.subprotocol == 'bar' server = await self.get_server(mock_wsconn, mock_request, client_subprotocols=['foo', 'bar'], server_subprotocols=['bar']) while server.connected: await asyncio.sleep(0.01) assert server.subprotocol == 'bar' server = await self.get_server(mock_wsconn, mock_request, client_subprotocols=['foo'], server_subprotocols=['foo', 'bar']) while server.connected: await asyncio.sleep(0.01) assert server.subprotocol == 'foo' server = await self.get_server(mock_wsconn, mock_request, client_subprotocols=['foo'], server_subprotocols=['bar', 'baz']) while server.connected: await asyncio.sleep(0.01) assert server.subprotocol is None server = await self.get_server(mock_wsconn, mock_request, client_subprotocols=['foo'], server_subprotocols=None) while server.connected: await asyncio.sleep(0.01) assert server.subprotocol is None simple-websocket-1.1.0/tests/test_server.py0000664000175000017500000002526014507561517016011 0ustar import time import unittest from unittest import mock import pytest # noqa: F401 from wsproto.events import Request, CloseConnection, TextMessage, \ BytesMessage, Ping, Pong import simple_websocket class SimpleWebSocketServerTestCase(unittest.TestCase): def get_server(self, mock_wsconn, environ, events=[], client_subprotocols=None, server_subprotocols=None, **kwargs): mock_wsconn().events.side_effect = \ [iter(ev) for ev in [[ Request(host='example.com', target='/ws', subprotocols=client_subprotocols or [])]] + events + [[CloseConnection(1000, 'bye')]]] mock_wsconn().send = lambda x: str(x).encode('utf-8') environ.update({ 'HTTP_HOST': 'example.com', 'HTTP_CONNECTION': 'Upgrade', 'HTTP_UPGRADE': 'websocket', 'HTTP_SEC_WEBSOCKET_KEY': 'Iv8io/9s+lYFgZWcXczP8Q==', 'HTTP_SEC_WEBSOCKET_VERSION': '13', }) return simple_websocket.Server.accept( environ, subprotocols=server_subprotocols, **kwargs) @mock.patch('simple_websocket.ws.WSConnection') def test_werkzeug(self, mock_wsconn): mock_socket = mock.MagicMock() mock_socket.recv.return_value = b'x' server = self.get_server(mock_wsconn, { 'werkzeug.socket': mock_socket, }) assert server.sock == mock_socket assert server.mode == 'werkzeug' assert server.receive_bytes == 4096 assert server.input_buffer == [] assert server.event.__class__.__name__ == 'Event' mock_wsconn().receive_data.assert_any_call( b'GET / HTTP/1.1\r\n' b'Host: example.com\r\n' b'Connection: Upgrade\r\n' b'Upgrade: websocket\r\n' b'Sec-Websocket-Key: Iv8io/9s+lYFgZWcXczP8Q==\r\n' b'Sec-Websocket-Version: 13\r\n\r\n') assert server.is_server @mock.patch('simple_websocket.ws.WSConnection') def test_gunicorn(self, mock_wsconn): mock_socket = mock.MagicMock() mock_socket.recv.return_value = b'x' server = self.get_server(mock_wsconn, { 'gunicorn.socket': mock_socket, }) assert server.sock == mock_socket assert server.mode == 'gunicorn' assert server.receive_bytes == 4096 assert server.input_buffer == [] assert server.event.__class__.__name__ == 'Event' mock_wsconn().receive_data.assert_any_call( b'GET / HTTP/1.1\r\n' b'Host: example.com\r\n' b'Connection: Upgrade\r\n' b'Upgrade: websocket\r\n' b'Sec-Websocket-Key: Iv8io/9s+lYFgZWcXczP8Q==\r\n' b'Sec-Websocket-Version: 13\r\n\r\n') assert server.is_server def test_no_socket(self): with pytest.raises(RuntimeError): self.get_server(mock.MagicMock(), {}) @mock.patch('simple_websocket.ws.WSConnection') def test_send(self, mock_wsconn): mock_socket = mock.MagicMock() mock_socket.recv.return_value = b'x' server = self.get_server(mock_wsconn, { 'werkzeug.socket': mock_socket, }) while server.connected: time.sleep(0.01) with pytest.raises(simple_websocket.ConnectionClosed): server.send('hello') server.connected = True server.send('hello') mock_socket.send.assert_called_with( b"TextMessage(data='hello', frame_finished=True, " b"message_finished=True)") server.connected = True server.send(b'hello') mock_socket.send.assert_called_with( b"Message(data=b'hello', frame_finished=True, " b"message_finished=True)") @mock.patch('simple_websocket.ws.WSConnection') def test_receive(self, mock_wsconn): mock_socket = mock.MagicMock() mock_socket.recv.return_value = b'x' server = self.get_server(mock_wsconn, { 'werkzeug.socket': mock_socket, }, events=[ [TextMessage('hello')], [BytesMessage(b'hello')], ]) while server.connected: time.sleep(0.01) server.connected = True assert server.receive() == 'hello' assert server.receive() == b'hello' assert server.receive(timeout=0) is None @mock.patch('simple_websocket.ws.WSConnection') def test_receive_after_close(self, mock_wsconn): mock_socket = mock.MagicMock() mock_socket.recv.return_value = b'x' server = self.get_server(mock_wsconn, { 'werkzeug.socket': mock_socket, }, events=[ [TextMessage('hello')], ]) while server.connected: time.sleep(0.01) assert server.receive() == 'hello' with pytest.raises(simple_websocket.ConnectionClosed): server.receive() @mock.patch('simple_websocket.ws.WSConnection') def test_receive_split_messages(self, mock_wsconn): mock_socket = mock.MagicMock() mock_socket.recv.return_value = b'x' server = self.get_server(mock_wsconn, { 'werkzeug.socket': mock_socket, }, events=[ [TextMessage('hel', message_finished=False)], [TextMessage('lo')], [TextMessage('he', message_finished=False)], [TextMessage('l', message_finished=False)], [TextMessage('lo')], [BytesMessage(b'hel', message_finished=False)], [BytesMessage(b'lo')], [BytesMessage(b'he', message_finished=False)], [BytesMessage(b'l', message_finished=False)], [BytesMessage(b'lo')], ]) while server.connected: time.sleep(0.01) server.connected = True assert server.receive() == 'hello' assert server.receive() == 'hello' assert server.receive() == b'hello' assert server.receive() == b'hello' assert server.receive(timeout=0) is None @mock.patch('simple_websocket.ws.WSConnection') def test_receive_ping(self, mock_wsconn): mock_socket = mock.MagicMock() mock_socket.recv.return_value = b'x' server = self.get_server(mock_wsconn, { 'werkzeug.socket': mock_socket, }, events=[ [Ping(b'hello')], ]) while server.connected: time.sleep(0.01) mock_socket.send.assert_any_call(b"Pong(payload=b'hello')") @mock.patch('simple_websocket.ws.WSConnection') def test_receive_empty(self, mock_wsconn): mock_socket = mock.MagicMock() mock_socket.recv.side_effect = [b'x', b'x', b''] server = self.get_server(mock_wsconn, { 'werkzeug.socket': mock_socket, }, events=[ [TextMessage('hello')], ]) while server.connected: time.sleep(0.01) server.connected = True assert server.receive() == 'hello' assert server.receive(timeout=0) is None @mock.patch('simple_websocket.ws.WSConnection') def test_receive_large(self, mock_wsconn): mock_socket = mock.MagicMock() mock_socket.recv.return_value = b'x' server = self.get_server(mock_wsconn, { 'werkzeug.socket': mock_socket, }, events=[ [TextMessage('hello')], [TextMessage('hello1')], ], max_message_size=5) while server.connected: time.sleep(0.01) server.connected = True assert server.receive() == 'hello' assert server.receive(timeout=0) is None @mock.patch('simple_websocket.ws.WSConnection') def test_close(self, mock_wsconn): mock_socket = mock.MagicMock() mock_socket.recv.return_value = b'x' server = self.get_server(mock_wsconn, { 'werkzeug.socket': mock_socket, }) while server.connected: time.sleep(0.01) with pytest.raises(simple_websocket.ConnectionClosed) as exc: server.close() assert str(exc.value) == 'Connection closed: 1000 bye' server.connected = True server.close() assert not server.connected mock_socket.send.assert_called_with( b'CloseConnection(code=, ' b'reason=None)') @mock.patch('simple_websocket.ws.WSConnection') @mock.patch('simple_websocket.ws.time') def test_ping_pong(self, mock_time, mock_wsconn): mock_sel = mock.MagicMock() mock_sel().select.side_effect = [True, True, False, False] mock_time.side_effect = [0, 1, 25.01, 25.02, 28, 52, 76] mock_socket = mock.MagicMock() mock_socket.recv.side_effect = [b'x', b'x'] server = self.get_server( mock_wsconn, {'werkzeug.socket': mock_socket}, events=[ [TextMessage('hello')], [Pong()], ], ping_interval=25, thread_class=mock.MagicMock(), selector_class=mock_sel) server._thread() assert mock_socket.send.call_count == 4 assert mock_socket.send.call_args_list[1][0][0].startswith(b'Ping') assert mock_socket.send.call_args_list[2][0][0].startswith(b'Ping') assert mock_socket.send.call_args_list[3][0][0].startswith(b'Close') @mock.patch('simple_websocket.ws.WSConnection') def test_subprotocols(self, mock_wsconn): mock_socket = mock.MagicMock() mock_socket.recv.return_value = b'x' server = self.get_server(mock_wsconn, { 'werkzeug.socket': mock_socket, }, client_subprotocols=['foo', 'bar'], server_subprotocols='bar') while server.connected: time.sleep(0.01) assert server.subprotocol == 'bar' server = self.get_server(mock_wsconn, { 'werkzeug.socket': mock_socket, }, client_subprotocols=['foo', 'bar'], server_subprotocols=['bar']) while server.connected: time.sleep(0.01) assert server.subprotocol == 'bar' server = self.get_server(mock_wsconn, { 'werkzeug.socket': mock_socket, }, client_subprotocols=['foo'], server_subprotocols=['foo', 'bar']) while server.connected: time.sleep(0.01) assert server.subprotocol == 'foo' server = self.get_server(mock_wsconn, { 'werkzeug.socket': mock_socket, }, client_subprotocols=['foo'], server_subprotocols=['bar', 'baz']) while server.connected: time.sleep(0.01) assert server.subprotocol is None server = self.get_server(mock_wsconn, { 'werkzeug.socket': mock_socket, }, client_subprotocols=['foo'], server_subprotocols=None) while server.connected: time.sleep(0.01) assert server.subprotocol is None simple-websocket-1.1.0/tests/__init__.py0000664000175000017500000000000014444310476015160 0ustar simple-websocket-1.1.0/tests/test_client.py0000664000175000017500000001664414702050766015763 0ustar import time import unittest from unittest import mock import pytest # noqa: F401 from wsproto.events import AcceptConnection, CloseConnection, TextMessage, \ BytesMessage, Ping import simple_websocket class SimpleWebSocketClientTestCase(unittest.TestCase): def get_client(self, mock_wsconn, url, events=[], subprotocols=None, headers=None): mock_wsconn().events.side_effect = \ [iter(ev) for ev in [[AcceptConnection()]] + events + [[CloseConnection(1000)]]] mock_wsconn().send = lambda x: str(x).encode('utf-8') return simple_websocket.Client.connect(url, subprotocols=subprotocols, headers=headers) @mock.patch('simple_websocket.ws.socket.socket') @mock.patch('simple_websocket.ws.WSConnection') def test_make_client(self, mock_wsconn, mock_socket): mock_socket.return_value.recv.return_value = b'x' client = self.get_client(mock_wsconn, 'ws://example.com/ws?a=1') assert client.sock == mock_socket() assert client.receive_bytes == 4096 assert client.input_buffer == [] assert client.event.__class__.__name__ == 'Event' client.sock.send.assert_called_with( b"Request(host='example.com', target='/ws?a=1', extensions=[], " b"extra_headers=[], subprotocols=[])") assert not client.is_server assert client.host == 'example.com' assert client.port == 80 assert client.path == '/ws?a=1' @mock.patch('simple_websocket.ws.socket.socket') @mock.patch('simple_websocket.ws.WSConnection') def test_make_client_subprotocol(self, mock_wsconn, mock_socket): mock_socket.return_value.recv.return_value = b'x' client = self.get_client(mock_wsconn, 'ws://example.com/ws?a=1', subprotocols='foo') assert client.subprotocols == ['foo'] client.sock.send.assert_called_with( b"Request(host='example.com', target='/ws?a=1', extensions=[], " b"extra_headers=[], subprotocols=['foo'])") @mock.patch('simple_websocket.ws.socket.socket') @mock.patch('simple_websocket.ws.WSConnection') def test_make_client_subprotocols(self, mock_wsconn, mock_socket): mock_socket.return_value.recv.return_value = b'x' client = self.get_client(mock_wsconn, 'ws://example.com/ws?a=1', subprotocols=['foo', 'bar']) assert client.subprotocols == ['foo', 'bar'] client.sock.send.assert_called_with( b"Request(host='example.com', target='/ws?a=1', extensions=[], " b"extra_headers=[], subprotocols=['foo', 'bar'])") @mock.patch('simple_websocket.ws.socket.socket') @mock.patch('simple_websocket.ws.WSConnection') def test_make_client_headers(self, mock_wsconn, mock_socket): mock_socket.return_value.recv.return_value = b'x' client = self.get_client(mock_wsconn, 'ws://example.com/ws?a=1', headers={'Foo': 'Bar'}) client.sock.send.assert_called_with( b"Request(host='example.com', target='/ws?a=1', extensions=[], " b"extra_headers=[('Foo', 'Bar')], subprotocols=[])") @mock.patch('simple_websocket.ws.socket.socket') @mock.patch('simple_websocket.ws.WSConnection') def test_make_client_headers2(self, mock_wsconn, mock_socket): mock_socket.return_value.recv.return_value = b'x' client = self.get_client(mock_wsconn, 'ws://example.com/ws?a=1', headers=[('Foo', 'Bar'), ('Foo', 'Baz')]) client.sock.send.assert_called_with( b"Request(host='example.com', target='/ws?a=1', extensions=[], " b"extra_headers=[('Foo', 'Bar'), ('Foo', 'Baz')], " b"subprotocols=[])") @mock.patch('simple_websocket.ws.socket.socket') @mock.patch('simple_websocket.ws.WSConnection') def test_send(self, mock_wsconn, mock_socket): mock_socket.return_value.recv.return_value = b'x' client = self.get_client(mock_wsconn, 'ws://example.com/ws') while client.connected: time.sleep(0.01) with pytest.raises(simple_websocket.ConnectionClosed): client.send('hello') client.connected = True client.send('hello') mock_socket().send.assert_called_with( b"TextMessage(data='hello', frame_finished=True, " b"message_finished=True)") client.connected = True client.send(b'hello') mock_socket().send.assert_called_with( b"Message(data=b'hello', frame_finished=True, " b"message_finished=True)") @mock.patch('simple_websocket.ws.socket.socket') @mock.patch('simple_websocket.ws.WSConnection') def test_receive(self, mock_wsconn, mock_socket): mock_socket.return_value.recv.return_value = b'x' client = self.get_client(mock_wsconn, 'ws://example.com/ws', events=[ [TextMessage('hello')], [BytesMessage(b'hello')], ]) while client.connected: time.sleep(0.01) client.connected = True assert client.receive() == 'hello' assert client.receive() == b'hello' assert client.receive(timeout=0) is None @mock.patch('simple_websocket.ws.socket.socket') @mock.patch('simple_websocket.ws.WSConnection') def test_receive_after_close(self, mock_wsconn, mock_socket): mock_socket.return_value.recv.return_value = b'x' client = self.get_client(mock_wsconn, 'ws://example.com/ws', events=[ [TextMessage('hello')], ]) while client.connected: time.sleep(0.01) assert client.receive() == 'hello' with pytest.raises(simple_websocket.ConnectionClosed): client.receive() @mock.patch('simple_websocket.ws.socket.socket') @mock.patch('simple_websocket.ws.WSConnection') def test_receive_ping(self, mock_wsconn, mock_socket): mock_socket.return_value.recv.return_value = b'x' client = self.get_client(mock_wsconn, 'ws://example.com/ws', events=[ [Ping(b'hello')], ]) while client.connected: time.sleep(0.01) mock_socket().send.assert_any_call(b"Pong(payload=b'hello')") @mock.patch('simple_websocket.ws.socket.socket') @mock.patch('simple_websocket.ws.WSConnection') def test_receive_empty(self, mock_wsconn, mock_socket): mock_socket.return_value.recv.side_effect = [b'x', b'x', b''] client = self.get_client(mock_wsconn, 'ws://example.com/ws', events=[ [TextMessage('hello')], ]) while client.connected: time.sleep(0.01) client.connected = True assert client.receive() == 'hello' assert client.receive(timeout=0) is None @mock.patch('simple_websocket.ws.socket.socket') @mock.patch('simple_websocket.ws.WSConnection') def test_close(self, mock_wsconn, mock_socket): mock_socket.return_value.recv.return_value = b'x' client = self.get_client(mock_wsconn, 'ws://example.com/ws') while client.connected: time.sleep(0.01) with pytest.raises(simple_websocket.ConnectionClosed): client.close() client.connected = True client.close() assert not client.connected mock_socket().send.assert_called_with( b'CloseConnection(code=, ' b'reason=None)') simple-websocket-1.1.0/tests/test_aioclient.py0000664000175000017500000002302614507561517016450 0ustar import asyncio import unittest from unittest import mock import pytest # noqa: F401 from wsproto.events import AcceptConnection, CloseConnection, TextMessage, \ BytesMessage, Ping import simple_websocket from .helpers import make_sync, AsyncMock class AioSimpleWebSocketClientTestCase(unittest.TestCase): async def get_client(self, mock_wsconn, url, events=[], subprotocols=None, headers=None): mock_wsconn().events.side_effect = \ [iter(ev) for ev in [[AcceptConnection()]] + events + [[CloseConnection(1000)]]] mock_wsconn().send = lambda x: str(x).encode('utf-8') return await simple_websocket.AioClient.connect( url, subprotocols=subprotocols, headers=headers) @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_make_client(self, mock_wsconn, mock_open_connection): rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) client = await self.get_client(mock_wsconn, 'ws://example.com/ws?a=1') assert client.rsock == rsock assert client.wsock == wsock assert client.receive_bytes == 4096 assert client.input_buffer == [] assert client.event.__class__.__name__ == 'Event' client.wsock.write.assert_called_with( b"Request(host='example.com', target='/ws?a=1', extensions=[], " b"extra_headers=[], subprotocols=[])") assert not client.is_server assert client.host == 'example.com' assert client.port == 80 assert client.path == '/ws?a=1' @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_make_client_subprotocol(self, mock_wsconn, mock_open_connection): rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) client = await self.get_client(mock_wsconn, 'ws://example.com/ws?a=1', subprotocols='foo') assert client.subprotocols == ['foo'] client.wsock.write.assert_called_with( b"Request(host='example.com', target='/ws?a=1', extensions=[], " b"extra_headers=[], subprotocols=['foo'])") @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_make_client_subprotocols(self, mock_wsconn, mock_open_connection): rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) client = await self.get_client(mock_wsconn, 'ws://example.com/ws?a=1', subprotocols=['foo', 'bar']) assert client.subprotocols == ['foo', 'bar'] client.wsock.write.assert_called_with( b"Request(host='example.com', target='/ws?a=1', extensions=[], " b"extra_headers=[], subprotocols=['foo', 'bar'])") @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_make_client_headers(self, mock_wsconn, mock_open_connection): rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) client = await self.get_client(mock_wsconn, 'ws://example.com/ws?a=1', headers={'Foo': 'Bar'}) client.wsock.write.assert_called_with( b"Request(host='example.com', target='/ws?a=1', extensions=[], " b"extra_headers=[('Foo', 'Bar')], subprotocols=[])") @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_make_client_headers2(self, mock_wsconn, mock_open_connection): rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) client = await self.get_client( mock_wsconn, 'ws://example.com/ws?a=1', headers=[('Foo', 'Bar'), ('Foo', 'Baz')]) client.wsock.write.assert_called_with( b"Request(host='example.com', target='/ws?a=1', extensions=[], " b"extra_headers=[('Foo', 'Bar'), ('Foo', 'Baz')], " b"subprotocols=[])") @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_send(self, mock_wsconn, mock_open_connection): rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) client = await self.get_client(mock_wsconn, 'ws://example.com/ws') while client.connected: await asyncio.sleep(0.01) with pytest.raises(simple_websocket.ConnectionClosed): await client.send('hello') client.connected = True await client.send('hello') wsock.write.assert_called_with( b"TextMessage(data='hello', frame_finished=True, " b"message_finished=True)") client.connected = True await client.send(b'hello') wsock.write.assert_called_with( b"Message(data=b'hello', frame_finished=True, " b"message_finished=True)") @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_receive(self, mock_wsconn, mock_open_connection): rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) client = await self.get_client( mock_wsconn, 'ws://example.com/ws', events=[ [TextMessage('hello')], [BytesMessage(b'hello')], ]) while client.connected: await asyncio.sleep(0.01) client.connected = True assert await client.receive() == 'hello' assert await client.receive() == b'hello' assert await client.receive(timeout=0) is None @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_receive_after_close(self, mock_wsconn, mock_open_connection): rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) client = await self.get_client( mock_wsconn, 'ws://example.com/ws', events=[ [TextMessage('hello')], ]) while client.connected: await asyncio.sleep(0.01) assert await client.receive() == 'hello' with pytest.raises(simple_websocket.ConnectionClosed): await client.receive() @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_receive_ping(self, mock_wsconn, mock_open_connection): rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) client = await self.get_client( mock_wsconn, 'ws://example.com/ws', events=[ [Ping(b'hello')], ]) while client.connected: await asyncio.sleep(0.01) wsock.write.assert_any_call(b"Pong(payload=b'hello')") @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_receive_empty(self, mock_wsconn, mock_open_connection): rsock = mock.MagicMock(read=AsyncMock(side_effect=[b'x', b'x', b''])) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) client = await self.get_client( mock_wsconn, 'ws://example.com/ws', events=[ [TextMessage('hello')], ]) while client.connected: await asyncio.sleep(0.01) client.connected = True assert await client.receive() == 'hello' assert await client.receive(timeout=0) is None @make_sync @mock.patch('simple_websocket.aiows.asyncio.open_connection') @mock.patch('simple_websocket.aiows.WSConnection') async def test_close(self, mock_wsconn, mock_open_connection): rsock = mock.MagicMock(read=AsyncMock(return_value=b'x')) wsock = mock.MagicMock() mock_open_connection.return_value = (rsock, wsock) client = await self.get_client( mock_wsconn, 'ws://example.com/ws') while client.connected: await asyncio.sleep(0.01) with pytest.raises(simple_websocket.ConnectionClosed): await client.close() client.connected = True await client.close() assert not client.connected wsock.write.assert_called_with( b'CloseConnection(code=, ' b'reason=None)') simple-websocket-1.1.0/tests/helpers.py0000664000175000017500000000110614507561517015077 0ustar import asyncio from unittest import mock def AsyncMock(*args, **kwargs): """Return a mock asynchronous function.""" m = mock.MagicMock(*args, **kwargs) async def mock_coro(*args, **kwargs): return m(*args, **kwargs) mock_coro.mock = m return mock_coro def _run(coro): """Run the given coroutine.""" return asyncio.get_event_loop().run_until_complete(coro) def make_sync(coro): """Wrap a coroutine so that it can be executed by pytest.""" def wrapper(*args, **kwargs): return _run(coro(*args, **kwargs)) return wrapper simple-websocket-1.1.0/MANIFEST.in0000664000175000017500000000020614512750450013447 0ustar include README.md LICENSE tox.ini recursive-include docs * recursive-exclude docs/_build * recursive-include tests * exclude **/*.pyc