Python manager.WebSocketManager类代码示例

class WSGIServer(_WSGIServer):
    def initialize_websockets_manager(self):
        """
        Call thos to start the underlying websockets
        manager. Make sure to call it once your server
        is created.
        """
        self.manager = WebSocketManager()
        self.manager.start()

    def shutdown_request(self, request):
        """
        The base class would close our socket
        if we didn't override it.
        """
        pass

    def link_websocket_to_server(self, ws):
        """
        Call this from your WSGI handler when a websocket
        has been created.
        """
        self.manager.add(ws)

    def server_close(self):
        """
        Properly initiate closing handshakes on
        all websockets when the WSGI server terminates.
        """
        if hasattr(self, 'manager'):
            self.manager.close_all()
            self.manager.stop()
            self.manager.join()
            delattr(self, 'manager')
        _WSGIServer.server_close(self)

    def test_websocket_close_all(self, MockSelectPoller):
        m = WebSocketManager(poller=MockSelectPoller())

        ws = MagicMock()
        m.add(ws)
        m.close_all()
        ws.terminate.assert_call_once_with(1001, 'Server is shutting down')

    def test_websocket_close_all(self, MockSelectPoller):
        m = WebSocketManager(poller=MockSelectPoller())

        ws = MagicMock()
        m.add(ws)
        m.close_all()
        ws.close.assert_called_once_with(code=1001, reason='Server is shutting down')

    def test_broadcast(self, MockSelectPoller):
        m = WebSocketManager(poller=MockSelectPoller())

        ws = MagicMock()
        ws.terminated = False
        m.add(ws)

        m.broadcast(b'hello there')
        ws.send.assert_call_once_with(b'hello there')

    def test_mainloop_can_be_stopped_when_no_websocket_were_registered(self, MockSelectPoller):
        m = WebSocketManager(poller=MockSelectPoller())
        self.assertFalse(m.running)
        
        m.start()
        self.assertTrue(m.running)

        m.stop()
        self.assertFalse(m.running)

    def test_add_and_remove_websocket(self, MockSelectPoller):
        m = WebSocketManager(poller=MockSelectPoller())

        ws = MagicMock()
        ws.sock.fileno.return_value = 1
        
        m.add(ws)
        m.poller.register.assert_call_once_with(ws)

        m.remove(ws)
        m.poller.unregister.assert_call_once_with(ws)

    def test_cannot_add_websocket_more_than_once(self, MockSelectPoller):
        m = WebSocketManager(poller=MockSelectPoller())

        ws = MagicMock()
        ws.sock.fileno.return_value = 1
        
        m.add(ws)
        self.assertEqual(len(m), 1)
        
        m.add(ws)
        self.assertEqual(len(m), 1)

    def test_broadcast_failure_must_not_break_caller(self, MockSelectPoller):
        m = WebSocketManager(poller=MockSelectPoller())

        ws = MagicMock()
        ws.terminated = False
        ws.send.side_effect = RuntimeError
        m.add(ws)

        try:
                m.broadcast(b'hello there')
        except:
                self.fail("Broadcasting shouldn't have failed")

 def test_mainloop_can_be_stopped(self, MockSelectPoller):
     m = WebSocketManager(poller=MockSelectPoller())
     
     def poll():
         yield 1
         m.stop()
         yield 2
         
     m.poller.poll.return_value = poll()
     self.assertFalse(m.running)
     
     m.start()
     # just make sure it had the time to finish
     time.sleep(0.1)
     self.assertFalse(m.running)

 def __init__(self, path):
     """ path is directory where static files are stored
     """
     self.path = path
     self.wsapp = WebSocketWSGIApplication(handler_cls=EchoWebSocket)
     self.manager = WebSocketManager()
     self.manager.start()

 def initialize_websockets_manager(self):
     """
     Call thos to start the underlying websockets
     manager. Make sure to call it once your server
     is created.
     """
     self.manager = WebSocketManager()
     self.manager.start()

    def execute(self, commands, environment={}):
        """Execute a command on the container.

        In pylxd 2.2, this method will be renamed `execute` and the existing
        `execute` method removed.
        """
        if not _ws4py_installed:
            raise ValueError(
                'This feature requires the optional ws4py library.')
        if isinstance(commands, six.string_types):
            raise TypeError("First argument must be a list.")
        response = self.api['exec'].post(json={
            'command': commands,
            'environment': environment,
            'wait-for-websocket': True,
            'interactive': False,
        })

        fds = response.json()['metadata']['metadata']['fds']
        operation_id = response.json()['operation'].split('/')[-1]
        parsed = parse.urlparse(
            self.client.api.operations[operation_id].websocket._api_endpoint)

        manager = WebSocketManager()

        stdin = _StdinWebsocket(self.client.websocket_url)
        stdin.resource = '{}?secret={}'.format(parsed.path, fds['0'])
        stdin.connect()
        stdout = _CommandWebsocketClient(manager, self.client.websocket_url)
        stdout.resource = '{}?secret={}'.format(parsed.path, fds['1'])
        stdout.connect()
        stderr = _CommandWebsocketClient(manager, self.client.websocket_url)
        stderr.resource = '{}?secret={}'.format(parsed.path, fds['2'])
        stderr.connect()

        manager.start()

        while len(manager.websockets.values()) > 0:
            time.sleep(.1)

        operation = self.client.operations.get(operation_id)
        return _ContainerExecuteResult(
            operation.metadata['return'], stdout.data, stderr.data)

def run(script_options):
    global logger
    level  = logging.DEBUG if script_options.verbose else logging.INFO
    logger = configure_logger(level = level)

    mgr = WebSocketManager()

    try:
        mgr.start()
        clients = []

        # Connect
        for connection_idx in range(script_options.concurrency):
            client = EchoClient(script_options.url, mgr,
                                script_options.ca, script_options.key, script_options.cert)
            client.connect()
            clients.append(client)

        logger.info("%d clients are connected" % (connection_idx + 1))

        # Send
        msg = getMessage(script_options)
        if msg:
            msg = json.write(msg)
            logger.info("Sending messages (num=%d):\n%s", script_options.num, msg)
            for client in clients:
                for _ in range(script_options.num):
                    client.send(msg)
                    time.sleep(SEND_INTERVAL)
            logger.info("Done sending")

        # Sleep before disconnecting
        logger.info("Sleeping for %d s before disconnecting",
                    script_options.interval)
        time.sleep(script_options.interval)

    except KeyboardInterrupt:
        logger.info("Interrupted by user")
    finally:
        logger.info("Disconnecting!")
        mgr.close_all(code    = 1000,
                      message = "Client is closing the connection")
        mgr.stop()
        mgr.join()

    def execute(self, commands, environment={}):
        """Execute a command on the container."""
        if isinstance(commands, six.string_types):
            raise TypeError("First argument must be a list.")
        response = self.api['exec'].post(json={
            'command': commands,
            'environment': environment,
            'wait-for-websocket': True,
            'interactive': False,
        })

        fds = response.json()['metadata']['metadata']['fds']
        operation_id = response.json()['operation'].split('/')[-1]
        parsed = parse.urlparse(
            self.client.api.operations[operation_id].websocket._api_endpoint)

        manager = WebSocketManager()

        stdin = _StdinWebsocket(manager, self.client.websocket_url)
        stdin.resource = '{}?secret={}'.format(parsed.path, fds['0'])
        stdin.connect()
        stdout = _CommandWebsocketClient(manager, self.client.websocket_url)
        stdout.resource = '{}?secret={}'.format(parsed.path, fds['1'])
        stdout.connect()
        stderr = _CommandWebsocketClient(manager, self.client.websocket_url)
        stderr.resource = '{}?secret={}'.format(parsed.path, fds['2'])
        stderr.connect()

        manager.start()

        while True:  # pragma: no cover
            for websocket in manager.websockets.values():
                if not websocket.terminated:
                    break
            else:
                break
            time.sleep(1)

        return stdout.data, stderr.data

    def test_cannot_remove_unregistered_websocket(self, MockSelectPoller):
        m = WebSocketManager(poller=MockSelectPoller())

        ws = MagicMock()
        ws.sock.fileno.return_value = 1
        
        m.remove(ws)
        self.assertEqual(len(m), 0)
        self.assertFalse(m.poller.unregister.called)
        
        m.add(ws)
        self.assertEqual(len(m), 1)
        m.remove(ws)
        self.assertEqual(len(m), 0)
        m.poller.unregister.assert_call_once_with(ws)
        m.poller.reset_mock()
        
        m.remove(ws)
        self.assertEqual(len(m), 0)
        self.assertFalse(m.poller.unregister.called)