How to use ws_handler method in devtools-proxy

Best Python code snippet using devtools-proxy_python

cherrypyserver.py

Source:cherrypyserver.py

1# -*- coding: utf-8 -*-2__doc__ = """3WebSocket within CherryPy is a tricky bit since CherryPy is4a threaded server which would choke quickly if each thread5of the server were kept attached to a long living connection6that WebSocket expects.7In order to work around this constraint, we take some advantage8of some internals of CherryPy as well as the introspection9Python provides.10Basically, when the WebSocket handshake is complete, we take over11the socket and let CherryPy take back the thread that was12associated with the upgrade request.13These operations require a bit of work at various levels of14the CherryPy framework but this module takes care of them15and from your application's perspective, this is abstracted.16Here are the various utilities provided by this module:17 * WebSocketTool: The tool is in charge to perform the18                  HTTP upgrade and detach the socket from19                  CherryPy. It runs at various hook points of the20                  request's processing. Enable that tool at21                  any path you wish to handle as a WebSocket22                  handler.23 * WebSocketPlugin: The plugin tracks the instanciated web socket handlers.24                    It also cleans out websocket handler which connection25                    have been closed down. The websocket connection then26                    runs in its own thread that this plugin manages.27Simple usage example:28.. code-block:: python29    :linenos:30    import cherrypy31    from ambari_ws4py.server.cherrypyserver import WebSocketPlugin, WebSocketTool32    from ambari_ws4py.websocket import EchoWebSocket33    cherrypy.config.update({'server.socket_port': 9000})34    WebSocketPlugin(cherrypy.engine).subscribe()35    cherrypy.tools.websocket = WebSocketTool()36    class Root(object):37        @cherrypy.expose38        def index(self):39            return 'some HTML with a websocket javascript connection'40        @cherrypy.expose41        def ws(self):42            pass43    cherrypy.quickstart(Root(), '/', config={'/ws': {'tools.websocket.on': True,44                                                     'tools.websocket.handler_cls': EchoWebSocket}})45Note that you can set the handler class on per-path basis,46meaning you could also dynamically change the class based47on other envrionmental settings (is the user authenticated for ex).48"""49import base6450from hashlib import sha151import inspect52import threading53import cherrypy54from cherrypy import Tool55from cherrypy.process import plugins56try:57    from cheroot.server import HTTPConnection, HTTPRequest, KnownLengthRFile58except ImportError:59    from cherrypy.wsgiserver import HTTPConnection, HTTPRequest, KnownLengthRFile    60from ambari_ws4py import WS_KEY, WS_VERSION61from ambari_ws4py.exc import HandshakeError62from ambari_ws4py.websocket import WebSocket63from ambari_ws4py.compat import py3k, get_connection, detach_connection64from ambari_ws4py.manager import WebSocketManager65__all__ = ['WebSocketTool', 'WebSocketPlugin']66class WebSocketTool(Tool):67    def __init__(self):68        Tool.__init__(self, 'before_request_body', self.upgrade)69    def _setup(self):70        conf = self._merged_args()71        hooks = cherrypy.serving.request.hooks72        p = conf.pop("priority", getattr(self.callable, "priority",73                                         self._priority))74        hooks.attach(self._point, self.callable, priority=p, **conf)75        hooks.attach('before_finalize', self.complete,76                     priority=p)77        hooks.attach('on_end_resource', self.cleanup_headers,78                     priority=70)79        hooks.attach('on_end_request', self.start_handler,80                     priority=70)81    def upgrade(self, protocols=None, extensions=None, version=WS_VERSION,82                handler_cls=WebSocket, heartbeat_freq=None):83        """84        Performs the upgrade of the connection to the WebSocket85        protocol.86        The provided protocols may be a list of WebSocket87        protocols supported by the instance of the tool.88        When no list is provided and no protocol is either89        during the upgrade, then the protocol parameter is90        not taken into account. On the other hand,91        if the protocol from the handshake isn't part92        of the provided list, the upgrade fails immediatly.93        """94        request = cherrypy.serving.request95        request.process_request_body = False96        ws_protocols = None97        ws_location = None98        ws_version = version99        ws_key = None100        ws_extensions = []101        if request.method != 'GET':102            raise HandshakeError('HTTP method must be a GET')103        for key, expected_value in [('Upgrade', 'websocket'),104                                     ('Connection', 'upgrade')]:105            actual_value = request.headers.get(key, '').lower()106            if not actual_value:107                raise HandshakeError('Header %s is not defined' % key)108            if expected_value not in actual_value:109                raise HandshakeError('Illegal value for header %s: %s' %110                                     (key, actual_value))111        version = request.headers.get('Sec-WebSocket-Version')112        supported_versions = ', '.join([str(v) for v in ws_version])113        version_is_valid = False114        if version:115            try: version = int(version)116            except: pass117            else: version_is_valid = version in ws_version118        if not version_is_valid:119            cherrypy.response.headers['Sec-WebSocket-Version'] = supported_versions120            raise HandshakeError('Unhandled or missing WebSocket version')121        key = request.headers.get('Sec-WebSocket-Key')122        if key:123            ws_key = base64.b64decode(key.encode('utf-8'))124            if len(ws_key) != 16:125                raise HandshakeError("WebSocket key's length is invalid")126        protocols = protocols or []127        subprotocols = request.headers.get('Sec-WebSocket-Protocol')128        if subprotocols:129            ws_protocols = []130            for s in subprotocols.split(','):131                s = s.strip()132                if s in protocols:133                    ws_protocols.append(s)134        exts = extensions or []135        extensions = request.headers.get('Sec-WebSocket-Extensions')136        if extensions:137            for ext in extensions.split(','):138                ext = ext.strip()139                if ext in exts:140                    ws_extensions.append(ext)141        location = []142        include_port = False143        if request.scheme == "https":144            location.append("wss://")145            include_port = request.local.port != 443146        else:147            location.append("ws://")148            include_port = request.local.port != 80149        location.append('localhost')150        if include_port:151            location.append(":%d" % request.local.port)152        location.append(request.path_info)153        if request.query_string != "":154            location.append("?%s" % request.query_string)155        ws_location = ''.join(location)156        response = cherrypy.serving.response157        response.stream = True158        response.status = '101 Switching Protocols'159        response.headers['Content-Type'] = 'text/plain'160        response.headers['Upgrade'] = 'websocket'161        response.headers['Connection'] = 'Upgrade'162        response.headers['Sec-WebSocket-Version'] = str(version)163        response.headers['Sec-WebSocket-Accept'] = base64.b64encode(sha1(key.encode('utf-8') + WS_KEY).digest())164        if ws_protocols:165            response.headers['Sec-WebSocket-Protocol'] = ', '.join(ws_protocols)166        if ws_extensions:167            response.headers['Sec-WebSocket-Extensions'] = ','.join(ws_extensions)168        addr = (request.remote.ip, request.remote.port)169        rfile = request.rfile.rfile170        if isinstance(rfile, KnownLengthRFile):171            rfile = rfile.rfile172        ws_conn = get_connection(rfile)173        request.ws_handler = handler_cls(ws_conn, ws_protocols, ws_extensions,174                                         request.wsgi_environ.copy(),175                                         heartbeat_freq=heartbeat_freq)176    def complete(self):177        """178        Sets some internal flags of CherryPy so that it179        doesn't close the socket down.180        """181        self._set_internal_flags()182    def cleanup_headers(self):183        """184        Some clients aren't that smart when it comes to185        headers lookup.186        """187        response = cherrypy.response188        if not response.header_list:189            return190        headers = response.header_list[:]191        for (k, v) in headers:192            if k[:7] == 'Sec-Web':193                response.header_list.remove((k, v))194                response.header_list.append((k.replace('Sec-Websocket', 'Sec-WebSocket'), v))195    def start_handler(self):196        """197        Runs at the end of the request processing by calling198        the opened method of the handler.199        """200        request = cherrypy.request201        if not hasattr(request, 'ws_handler'):202            return203        addr = (request.remote.ip, request.remote.port)204        ws_handler = request.ws_handler205        request.ws_handler = None206        delattr(request, 'ws_handler')207        # By doing this we detach the socket from208        # the CherryPy stack avoiding memory leaks209        detach_connection(request.rfile.rfile)210        cherrypy.engine.publish('handle-websocket', ws_handler, addr)211    def _set_internal_flags(self):212        """213        CherryPy has two internal flags that we are interested in214        to enable WebSocket within the server. They can't be set via215        a public API and considering I'd want to make this extension216        as compatible as possible whilst refraining in exposing more217        than should be within CherryPy, I prefer performing a bit218        of introspection to set those flags. Even by Python standards219        such introspection isn't the cleanest but it works well220        enough in this case.221        This also means that we do that only on WebSocket222        connections rather than globally and therefore we do not223        harm the rest of the HTTP server.224        """225        current = inspect.currentframe()226        while True:227            if not current:228                break229            _locals = current.f_locals230            if 'self' in _locals:231                if isinstance(_locals['self'], HTTPRequest):232                    _locals['self'].close_connection = True233                if isinstance(_locals['self'], HTTPConnection):234                    _locals['self'].linger = True235                    # HTTPConnection is more inner than236                    # HTTPRequest so we can leave once237                    # we're done here238                    return239            _locals = None240            current = current.f_back241class WebSocketPlugin(plugins.SimplePlugin):242    def __init__(self, bus):243        plugins.SimplePlugin.__init__(self, bus)244        self.manager = WebSocketManager()245    def start(self):246        self.bus.log("Starting WebSocket processing")247        self.bus.subscribe('stop', self.cleanup)248        self.bus.subscribe('handle-websocket', self.handle)249        self.bus.subscribe('websocket-broadcast', self.broadcast)250        self.manager.start()251    def stop(self):252        self.bus.log("Terminating WebSocket processing")253        self.bus.unsubscribe('stop', self.cleanup)254        self.bus.unsubscribe('handle-websocket', self.handle)255        self.bus.unsubscribe('websocket-broadcast', self.broadcast)256    def handle(self, ws_handler, peer_addr):257        """258        Tracks the provided handler.259        :param ws_handler: websocket handler instance260        :param peer_addr: remote peer address for tracing purpose261        """262        self.manager.add(ws_handler)263    def cleanup(self):264        """265        Terminate all connections and clear the pool. Executed when the engine stops.266        """267        self.manager.close_all()268        self.manager.stop()269        self.manager.join()270    def broadcast(self, message, binary=False):271        """272        Broadcasts a message to all connected clients known to273        the server.274        :param message: a message suitable to pass to the send() method275          of the connected handler.276        :param binary: whether or not the message is a binary one277        """278        self.manager.broadcast(message, binary)279if __name__ == '__main__':280    import random281    cherrypy.config.update({'server.socket_host': '127.0.0.1',282                            'server.socket_port': 9000})283    WebSocketPlugin(cherrypy.engine).subscribe()284    cherrypy.tools.websocket = WebSocketTool()285    class Root(object):286        @cherrypy.expose287        @cherrypy.tools.websocket(on=False)288        def ws(self):289            return """<html>290        <head>291          <script type='application/javascript' src='https://ajax.googleapis.com/ajax/libs/jquery/1.8.3/jquery.min.js'> </script>292          <script type='application/javascript'>293            $(document).ready(function() {294              var ws = new WebSocket('ws://192.168.0.10:9000/');295              ws.onmessage = function (evt) {296                 $('#chat').val($('#chat').val() + evt.data + '\\n');297              };298              ws.onopen = function() {299                 ws.send("Hello there");300              };301              ws.onclose = function(evt) {302                $('#chat').val($('#chat').val() + 'Connection closed by server: ' + evt.code + ' \"' + evt.reason + '\"\\n');303              };304              $('#chatform').submit(function() {305                 ws.send('%(username)s: ' + $('#message').val());306                 $('#message').val("");307                 return false;308              });309            });310          </script>311        </head>312        <body>313        <form action='/echo' id='chatform' method='get'>314          <textarea id='chat' cols='35' rows='10'></textarea>315          <br />316          <label for='message'>%(username)s: </label><input type='text' id='message' />317          <input type='submit' value='Send' />318          </form>319        </body>320        </html>321        """ % {'username': "User%d" % random.randint(0, 100)}322        @cherrypy.expose323        def index(self):324            cherrypy.log("Handler created: %s" % repr(cherrypy.request.ws_handler))325    cherrypy.quickstart(Root(), '/', config={'/': {'tools.websocket.on': True,...

test_session.py

Source:test_session.py

1import asyncio2import pytest3from ..schema import SERVERS_RESPONSE4def assert_status_set(handler, expected_statuses, language_server=None):5    handler.get()6    payload = handler._payload7    errors = list(SERVERS_RESPONSE.iter_errors(payload))8    assert not errors9    statuses = {10        session["status"]11        for session_server, session in payload["sessions"].items()12        if language_server is None or language_server == session_server13    }14    assert statuses == expected_statuses, payload15@pytest.mark.asyncio16async def test_start_known(known_server, handlers, jsonrpc_init_msg):17    """ will a process start for a known server if a handler starts?18    """19    handler, ws_handler = handlers20    manager = handler.manager21    manager.initialize()22    assert_status_set(handler, {"not_started"})23    ws_handler.open(known_server)24    session = manager.sessions[ws_handler.language_server]25    assert session.process is not None26    assert_status_set(handler, {"started"}, known_server)27    await ws_handler.on_message(jsonrpc_init_msg)28    try:29        await asyncio.wait_for(ws_handler._messages_wrote.get(), 20)30        ws_handler._messages_wrote.task_done()31    finally:32        ws_handler.on_close()33    assert not session.handlers34    assert not session.process35    assert_status_set(handler, {"stopped"}, known_server)36    assert_status_set(handler, {"stopped", "not_started"})37@pytest.mark.asyncio38async def test_start_unknown(known_unknown_server, handlers, jsonrpc_init_msg):39    """ will a process not start for an unknown server if a handler starts?40    """41    handler, ws_handler = handlers42    manager = handler.manager43    manager.initialize()44    assert_status_set(handler, {"not_started"})45    ws_handler.open(known_unknown_server)46    assert_status_set(handler, {"not_started"})47    await ws_handler.on_message(jsonrpc_init_msg)48    assert_status_set(handler, {"not_started"})49    ws_handler.on_close()50    assert not manager.sessions.get(ws_handler.language_server)...

Automation Testing Tutorials

Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.