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)

    def test_websocket_terminated_from_mainloop(self, MockSelectPoller):
        m = WebSocketManager(poller=MockSelectPoller())
        m.poller.poll.return_value = [1]

        ws = MagicMock()
        
        ws.terminated = False
        ws.sock.fileno.return_value = 1
        ws.once.return_value = False
        
        m.add(ws)
        m.start()
        
        ws.terminate.assert_call_once_with()
        
        m.stop()

class FileServer(object):
    """ Serves static files from a directory.
    """

    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 __call__(self, environ, start_response):
        """ WSGI entry point
        """
        # Upgrade header means websockets...
        upgrade_header = environ.get('HTTP_UPGRADE', '').lower()
        if upgrade_header:
            environ['ws4py.socket'] = get_connection(environ['wsgi.input'])
            # This will make a websocket, hopefully!
            ret = self.wsapp(environ, start_response)
            if 'ws4py.websocket' in environ:
                self.manager.add(environ.pop('ws4py.websocket'))
            return ret

        # Find path to file to server
        path_info = environ["PATH_INFO"]

        if not path_info or path_info == "/":
            path_info = "/index.html"

        file_path = os.path.join(self.path, path_info[1:])

        # If file does not exist, return 404
        if not os.path.exists(file_path):
            return self._not_found(start_response)

        # Guess mimetype of file based on file extension
        mimetype = mimetypes.guess_type(file_path)[0]

        # If we can't guess mimetype, return a 403 Forbidden
        if mimetype is None:
            return self._forbidden(start_response)

        # Get size of file
        size = os.path.getsize(file_path)

        # Create headers and start response
        headers = [
            ("Content-type", mimetype),
            ("Content-length", str(size)),
        ]

        start_response("200 OK", headers)

        # Send file
        return self._send_file(file_path, size)

    def _send_file(self, file_path, size):
        """ A generator function which returns the blocks in a file, one at
        a time.

        """
        with open(file_path) as f:
            block = f.read(BLOCK_SIZE)
            while block:
                yield block
                block = f.read(BLOCK_SIZE)

    def _not_found(self, start_response):
        start_response("404 NOT FOUND", [("Content-type", "text/plain")])
        return ["Not found", ]

    def _forbidden(self, start_response):
        start_response("403 FORBIDDEN", [("Content-type", "text/plain")])
        return ["Forbidden", ]

class WebSocketPlugin(plugins.SimplePlugin):
    def __init__(self, bus):
        plugins.SimplePlugin.__init__(self, bus)
        self.manager = WebSocketManager()

    def start(self):
        self.bus.log("Starting WebSocket processing")
        self.bus.subscribe('stop', self.cleanup)
        self.bus.subscribe('handle-websocket', self.handle)
        self.bus.subscribe('websocket-broadcast', self.broadcast)
        self.manager.start()

    def stop(self):
        self.bus.log("Terminating WebSocket processing")
        self.bus.unsubscribe('stop', self.cleanup)
        self.bus.unsubscribe('handle-websocket', self.handle)
        self.bus.unsubscribe('websocket-broadcast', self.broadcast)

    def handle(self, ws_handler, peer_addr):
        """
        Tracks the provided handler.

        :param ws_handler: websocket handler instance
        :param peer_addr: remote peer address for tracing purpose
        """
        self.manager.add(ws_handler)

    def cleanup(self):
        """
        Terminate all connections and clear the pool. Executed when the engine stops.
        """
        self.manager.close_all()
        self.manager.stop()
        self.manager.join()

    def broadcast(self, message, binary=False):
        """
        Broadcasts a message to all connected clients known to
        the server.

        :param message: a message suitable to pass to the send() method
          of the connected handler.
        :param binary: whether or not the message is a binary one
        """
        self.manager.broadcast(message, binary)

 def __init__(self,start, poller=None):
     threading.Thread.__init__(self)
     WebSocketManager.__init__(self)
     self.timeout = start + (60*30)

 def __init__(self, bus):
     plugins.SimplePlugin.__init__(self, bus)
     self.manager = WebSocketManager()