How to use ws4py - 10 common examples

import time
import json
import multiprocessing

import pytest
from tornado import ioloop
from ws4py.client.tornadoclient import TornadoWebSocketClient

from omnibus.daemon import spawn_omnibusd
from omnibus.pubsub import PubSub

from .helpers import AssertEventuallyTest


class WebsocketClient(TornadoWebSocketClient):
    def opened(self):
        print('opened, send authenticate')
        self.send('!authenticate:123456')
        time.sleep(.1)

    def received_message(self, message):
        print('received message', message, message.data)

        type, payload = message.data.split(':', 1)

        if type == '!authenticate' and json.loads(payload).get('success', False):
            self.send('!subscribe:channel')
        elif type == '!subscribe' and json.loads(payload).get('success', False):
            self.send('channel:{"type": "type", "sender": null, "payload": {"test": "works!"}}')
            time.sleep(.2)
            self.send('channel:{"type": "type", "sender": null, "payload": {"test": "works2!"}}')

def test_sending_ping(self):
        tm = PingControlMessage("hello").single(mask=False)
        
        m = MagicMock()
        ws = WebSocket(sock=m)
        ws.ping("hello")
        m.sendall.assert_called_once_with(tm)

def test_send_generator_without_masking(self):
        tm0 = b'hello'
        tm1 = b'world'
        
        def datasource():
            yield tm0
            yield tm1

        gen = datasource()
        
        m = MagicMock()
        ws = WebSocket(sock=m)
        ws.send(gen)
        self.assertEqual(m.sendall.call_count, 2)
        self.assertRaises(StopIteration, next, gen)

def test_terminate_with_closing(self):
        m = MagicMock()
        s = MagicMock()
        c = MagicMock()
        cc = MagicMock()
        
        ws = WebSocket(sock=m)
        with patch.multiple(ws, closed=c, close_connection=cc):
            ws.stream = s
            ws.stream.closing = CloseControlMessage(code=1000, reason='test closing')
            ws.terminate()
            self.assertTrue(ws.client_terminated)
            self.assertTrue(ws.server_terminated)
            self.assertTrue(ws.terminated)
            c.assert_called_once_with(1000, b'test closing')
            cc.assert_called_once_with()
            self.assertIsNone(ws.stream)
            self.assertIsNone(ws.environ)

def test_send_bytes_with_masking(self):
        tm = TextMessage(b'hello world').single(mask=True)
        
        m = MagicMock()
        ws = WebSocket(sock=m)
        ws.stream = MagicMock()
        ws.stream.always_mask = True
        ws.stream.text_message.return_value.single.return_value = tm
        
        ws.send(b'hello world')
        m.sendall.assert_called_once_with(tm)

def test_invalid_encoded_bytes_on_continuation(self):
        s = Stream()

        f = Frame(opcode=OPCODE_TEXT, body=b'hello',
                  fin=0, masking_key=os.urandom(4)).build()
        s.parser.send(f)
        next(s.parser)

        f = Frame(opcode=OPCODE_CONTINUATION, body=b'h\xc3llo',
                  fin=1, masking_key=os.urandom(4)).build()
        s.parser.send(f)
        next(s.parser)

        self.assertNotEqual(s.errors, [])
        self.assertIsInstance(s.errors[0], CloseControlMessage)
        self.assertEqual(s.errors[0].code, 1007)

def test_incremental_parsing_small_7_bit_length(self):
        bytes = Frame(opcode=OPCODE_TEXT, body=b'hello', fin=1).build()

        f = Frame()
        map_on_bytes(f.parser.send, bytes)
        self.assertTrue(f.masking_key is None)
        self.assertEqual(f.payload_length, 5)

msg = os.urandom(16)
        key = os.urandom(4)
        f = Frame(opcode=OPCODE_BINARY, body=msg, fin=0, masking_key=key).build()
        s = Stream()
        self.assertEqual(s.has_message, False)
        s.parser.send(f)
        self.assertEqual(s.has_message, False)

        for i in range(3):
            f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=0, masking_key=key).build()
            s.parser.send(f)
            self.assertEqual(s.has_message, False)
            self.assertEqual(s.message.completed, False)
            self.assertEqual(s.message.opcode, OPCODE_BINARY)

        f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=1, masking_key=key).build()
        s.parser.send(f)
        self.assertEqual(s.has_message, True)
        self.assertEqual(s.message.completed, True)
        self.assertEqual(s.message.opcode, OPCODE_BINARY)

def test_non_zero_nor_one_fin(self):
        f = Frame(opcode=OPCODE_TEXT,
                  body=b'', fin=2)
        self.assertRaises(ValueError, f.build)

def test_incremental_text_message_received(self):
        msg = b'hello there'
        f = Frame(opcode=OPCODE_TEXT, body=msg, fin=1, masking_key=os.urandom(4)).build()
        s = Stream()
        self.assertEqual(s.has_message, False)
        bytes = f
        for index, byte in enumerate(bytes):
            s.parser.send(bytes[index:index+1])
        self.assertEqual(s.has_message, True)