codekingpro/portable-devtools

1import uuid2 3from socketio import packet4from socketio.pubsub_manager import PubSubManager5from werkzeug.test import EnvironBuilder6 7 8class SocketIOTestClient(object):9    """10    This class is useful for testing a Flask-SocketIO server. It works in a11    similar way to the Flask Test Client, but adapted to the Socket.IO server.12 13    :param app: The Flask application instance.14    :param socketio: The application's ``SocketIO`` instance.15    :param namespace: The namespace for the client. If not provided, the client16                      connects to the server on the global namespace.17    :param query_string: A string with custom query string arguments.18    :param headers: A dictionary with custom HTTP headers.19    :param auth: Optional authentication data, given as a dictionary.20    :param flask_test_client: The instance of the Flask test client21                              currently in use. Passing the Flask test22                              client is optional, but is necessary if you23                              want the Flask user session and any other24                              cookies set in HTTP routes accessible from25                              Socket.IO events.26    """27    clients = {}28 29    def __init__(self, app, socketio, namespace=None, query_string=None,30                 headers=None, auth=None, flask_test_client=None):31        def _mock_send_packet(eio_sid, pkt):32            # make sure the packet can be encoded and decoded33            epkt = pkt.encode()34            if not isinstance(epkt, list):35                pkt = packet.Packet(encoded_packet=epkt)36            else:37                pkt = packet.Packet(encoded_packet=epkt[0])38                for att in epkt[1:]:39                    pkt.add_attachment(att)40            client = self.clients.get(eio_sid)41            if not client:42                return43            if pkt.packet_type == packet.EVENT or \44                    pkt.packet_type == packet.BINARY_EVENT:45                if pkt.data[0] == 'message' or pkt.data[0] == 'json':46                    client.queue.append({47                        'name': pkt.data[0],48                        'args': pkt.data[1],49                        'namespace': pkt.namespace or '/'})50                else:51                    client.queue.append({52                        'name': pkt.data[0],53                        'args': pkt.data[1:],54                        'namespace': pkt.namespace or '/'})55            elif pkt.packet_type == packet.ACK or \56                    pkt.packet_type == packet.BINARY_ACK:57                client.acks = {'args': pkt.data,58                               'namespace': pkt.namespace or '/'}59            elif pkt.packet_type in [packet.DISCONNECT, packet.CONNECT_ERROR]:60                client.connected[pkt.namespace or '/'] = False61 62        _current_packet = None63 64        def _mock_send_eio_packet(eio_sid, eio_pkt):65            nonlocal _current_packet66            if _current_packet is not None:67                _current_packet.add_attachment(eio_pkt.data)68                if _current_packet.attachment_count == \69                        len(_current_packet.attachments):70                    _mock_send_packet(eio_sid, _current_packet)71                    _current_packet = None72            else:73                pkt = packet.Packet(encoded_packet=eio_pkt.data)74                if pkt.attachment_count == 0:75                    _mock_send_packet(eio_sid, pkt)76                else:77                    _current_packet = pkt78 79        self.app = app80        self.flask_test_client = flask_test_client81        self.eio_sid = uuid.uuid4().hex82        self.clients[self.eio_sid] = self83        self.callback_counter = 084        self.socketio = socketio85        self.connected = {}86        self.queue = []87        self.acks = None88        socketio.server._send_packet = _mock_send_packet89        socketio.server._send_eio_packet = _mock_send_eio_packet90        socketio.server.environ[self.eio_sid] = {}91        socketio.server.async_handlers = False      # easier to test when92        socketio.server.eio.async_handlers = False  # events are sync93        if isinstance(socketio.server.manager, PubSubManager):94            raise RuntimeError('Test client cannot be used with a message '95                               'queue. Disable the queue on your test '96                               'configuration.')97        socketio.server.manager.initialize()98        self.connect(namespace=namespace, query_string=query_string,99                     headers=headers, auth=auth)100 101    def is_connected(self, namespace=None):102        """Check if a namespace is connected.103 104        :param namespace: The namespace to check. The global namespace is105                         assumed if this argument is not provided.106        """107        return self.connected.get(namespace or '/', False)108 109    def connect(self, namespace=None, query_string=None, headers=None,110                auth=None):111        """Connect the client.112 113        :param namespace: The namespace for the client. If not provided, the114                          client connects to the server on the global115                          namespace.116        :param query_string: A string with custom query string arguments.117        :param headers: A dictionary with custom HTTP headers.118        :param auth: Optional authentication data, given as a dictionary.119 120        Note that it is usually not necessary to explicitly call this method,121        since a connection is automatically established when an instance of122        this class is created. An example where it this method would be useful123        is when the application accepts multiple namespace connections.124        """125        url = '/socket.io'126        namespace = namespace or '/'127        if query_string:128            if query_string[0] != '?':129                query_string = '?' + query_string130            url += query_string131        environ = EnvironBuilder(url, headers=headers).get_environ()132        environ['flask.app'] = self.app133        if self.flask_test_client:134            # inject cookies from Flask135            if hasattr(self.flask_test_client, '_add_cookies_to_wsgi'):136                # flask >= 2.3137                self.flask_test_client._add_cookies_to_wsgi(environ)138            else:  # pragma: no cover139                # flask < 2.3140                self.flask_test_client.cookie_jar.inject_wsgi(environ)141        self.socketio.server._handle_eio_connect(self.eio_sid, environ)142        pkt = packet.Packet(packet.CONNECT, auth, namespace=namespace)143        self.socketio.server._handle_eio_message(self.eio_sid, pkt.encode())144        sid = self.socketio.server.manager.sid_from_eio_sid(self.eio_sid,145                                                            namespace)146        if sid:147            self.connected[namespace] = True148 149    def disconnect(self, namespace=None):150        """Disconnect the client.151 152        :param namespace: The namespace to disconnect. The global namespace is153                         assumed if this argument is not provided.154        """155        if not self.is_connected(namespace):156            raise RuntimeError('not connected')157        pkt = packet.Packet(packet.DISCONNECT, namespace=namespace)158        self.socketio.server._handle_eio_message(self.eio_sid, pkt.encode())159        del self.connected[namespace or '/']160 161    def emit(self, event, *args, **kwargs):162        """Emit an event to the server.163 164        :param event: The event name.165        :param *args: The event arguments.166        :param callback: ``True`` if the client requests a callback, ``False``167                         if not. Note that client-side callbacks are not168                         implemented, a callback request will just tell the169                         server to provide the arguments to invoke the170                         callback, but no callback is invoked. Instead, the171                         arguments that the server provided for the callback172                         are returned by this function.173        :param namespace: The namespace of the event. The global namespace is174                          assumed if this argument is not provided.175        """176        namespace = kwargs.pop('namespace', None)177        if not self.is_connected(namespace):178            raise RuntimeError('not connected')179        callback = kwargs.pop('callback', False)180        id = None181        if callback:182            self.callback_counter += 1183            id = self.callback_counter184        pkt = packet.Packet(packet.EVENT, data=[event] + list(args),185                            namespace=namespace, id=id)186        encoded_pkt = pkt.encode()187        if isinstance(encoded_pkt, list):188            for epkt in encoded_pkt:189                self.socketio.server._handle_eio_message(self.eio_sid, epkt)190        else:191            self.socketio.server._handle_eio_message(self.eio_sid, encoded_pkt)192        if self.acks is not None:193            ack = self.acks194            self.acks = None195            return ack['args'][0] if len(ack['args']) == 1 \196                else ack['args']197 198    def send(self, data, json=False, callback=False, namespace=None):199        """Send a text or JSON message to the server.200 201        :param data: A string, dictionary or list to send to the server.202        :param json: ``True`` to send a JSON message, ``False`` to send a text203                     message.204        :param callback: ``True`` if the client requests a callback, ``False``205                         if not. Note that client-side callbacks are not206                         implemented, a callback request will just tell the207                         server to provide the arguments to invoke the208                         callback, but no callback is invoked. Instead, the209                         arguments that the server provided for the callback210                         are returned by this function.211        :param namespace: The namespace of the event. The global namespace is212                          assumed if this argument is not provided.213        """214        if json:215            msg = 'json'216        else:217            msg = 'message'218        return self.emit(msg, data, callback=callback, namespace=namespace)219 220    def get_received(self, namespace=None):221        """Return the list of messages received from the server.222 223        Since this is not a real client, any time the server emits an event,224        the event is simply stored. The test code can invoke this method to225        obtain the list of events that were received since the last call.226 227        :param namespace: The namespace to get events from. The global228                          namespace is assumed if this argument is not229                          provided.230        """231        if not self.is_connected(namespace):232            raise RuntimeError('not connected')233        namespace = namespace or '/'234        r = [pkt for pkt in self.queue if pkt['namespace'] == namespace]235        self.queue = [pkt for pkt in self.queue if pkt not in r]236        return r237