/tests/web/test_websockets.py

  1#!/usr/bin/env python
  2
  3
  4from __future__ import print_function
  5
  6
  7from circuits import Component
  8from circuits.web.servers import Server
  9from circuits.web.controllers import Controller
 10from circuits.net.sockets import close, write
 11from circuits.web.websockets import WebSocketClient, WebSocketsDispatcher
 12
 13
 14from .helpers import urlopen
 15
 16
 17class Echo(Component):
 18
 19    channel = "wsserver"
 20
 21    def init(self):
 22        self.clients = []
 23
 24    def connect(self, sock, host, port):
 25        self.clients.append(sock)
 26        print("WebSocket Client Connected:", host, port)
 27        self.fire(write(sock, "Welcome {0:s}:{1:d}".format(host, port)))
 28
 29    def disconnect(self, sock):
 30        self.clients.remove(sock)
 31
 32    def read(self, sock, data):
 33        self.fire(write(sock, "Received: " + data))
 34
 35
 36class Root(Controller):
 37
 38    def index(self):
 39        return "Hello World!"
 40
 41
 42class Client(Component):
 43
 44    channel = "ws"
 45
 46    def init(self, *args, **kwargs):
 47        self.response = None
 48
 49    def read(self, data):
 50        self.response = data
 51
 52
 53def test(manager, watcher, webapp):
 54    server = Server(0).register(manager)
 55    watcher.wait("ready")
 56
 57    echo = Echo().register(server)
 58    Root().register(server)
 59    watcher.wait("registered", channel="wsserver")
 60
 61    f = urlopen(webapp.server.http.base)
 62    s = f.read()
 63    assert s == b"Hello World!"
 64
 65    watcher.clear()
 66
 67    WebSocketsDispatcher("/websocket").register(server)
 68    watcher.wait("registered", channel="web")
 69
 70    uri = "ws://{0:s}:{1:d}/websocket".format(server.host, server.port)
 71
 72    WebSocketClient(uri).register(manager)
 73    client = Client().register(manager)
 74    watcher.wait("registered", channel="wsclient")
 75    watcher.wait("connected", channel="wsclient")
 76
 77    assert len(echo.clients) == 1
 78
 79    watcher.wait("read", channel="ws")
 80    assert client.response.startswith("Welcome")
 81    watcher.clear()
 82
 83    client.fire(write("Hello!"), "ws")
 84    watcher.wait("read", channel="ws")
 85    assert client.response == "Received: Hello!"
 86
 87    f = urlopen(webapp.server.http.base)
 88    s = f.read()
 89    assert s == b"Hello World!"
 90
 91    assert len(echo.clients) == 1
 92
 93    client.fire(close(), "ws")
 94    watcher.wait("disconnect", channel="wsserver")
 95    assert len(echo.clients) == 0
 96
 97    client.unregister()
 98    watcher.wait("unregistered")
 99    watcher.clear()
100
101    server.unregister()
102    watcher.wait("unregistered")