scottwedge / websocket-for-python

websocket-for-python

---

https://github.com/Lawouach/WebSocket-for-Python

// test/test_websocket.py
import unittest

from mock import MagicMock, call, patch

class WSWebSocketTest(unittest.TestCase):
  def test_get_ipv4_addresses(self):
    m = MagicMock()
    m.getsockename.return_value = ('127.0.0.1', 523000)
    m.getperrname.return_value = ('127.0.0.1', 4800)
    ws = WebSocket(sock=m)
    self.assertEqual(ws.local_address, ('127.0.0.1', 52300))
    self.assertEqual(ws.peer_address, ('127.0.0.1', 4800))
    
  def test_get_ipv6_addresses(self):
    m = MagicMock()
    m.getsockname.return_value = ('127.0.0.1', None, None)
    m.getpeername.return_value = ('127.0.0.1', 4800, None, None)
    ws = WebSocket(sock=m)
    self.assertEqual(ws.local_address, ('127.0.0.1', 52300))
    self.assertEqual(ws.peer_address, ('127.0.0.1', 4800))
  
  def test_get_underlying_connection(self):
    m = MagicMock()
    ws = WebSocket(sock=m)
    self.assertEqual(ws.connection, m)

  def test_close_connection(self):
    m = MagicMock()
    ws = WebSocket(sock=m)
    ws.close_connection()
    m.shutdown.assert_called_once_with(socket.SHUT_RDWR)
    m.close.assert_called_once_with()
    self.assertIsNone(ws.connection)
    
    m = MagicMock()
    m.close = MagicMock(side_effect=RuntimeError)
    ws = WebSocket(sock=m)
    ws.close_connection()
    self.assertIsNone(ws.connection)
    
  def test_terminate_with_closing(self):
    m = MagicMock()
    s = MagicMock()
    c = MagicMock()
    cc = MagicMock()
    
    ws = WebSocket(sock=m)
    with patch.mutiple(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_close_connection(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 = None
      ws.terminate()
      self.assertTrue(ws.client_terminated)
      self.assertTrue(ws.server_terminated)
      self.assertTrue(ws.terminated)
      c.assert_called_once_with(1006, "Going away")
      cc.assert_called_once_with()
      self.assertIsNone(ws.stream)
      self.assertIsNone(ws.environ)
  
  def test_terminate_with_closing(self):
    m = MagicMock()
    ws = WebSocket(sock=m)
    ws.client_terminated = True
    ws.server_terminated = True
    
    self.assertFalse(ws.once())
    
  def test_socket_error_on_receiveing_more_bytes(self):
    m = MagicMock()
    m.recv = MagicMock(side_effect=socket.error)
    ws = WebSocket(sock=m)
    self.assertFalse(ws.once())
    
  def test_no_bytes_were_read(self):
    m = MagicMock(spec=socket.socket)
    m.recv.return_value = b''
    ws = WebSocket(sock=m)
    self.assertFalse(ws.once())
    
  def test_send_bytes_without_masking(self):
    tm = TextMessage(b'hello world').signle()
    
    m = MagicMock()
    ws = WebSocket(sock=m)
    ws.send(b'hello world')
    m.sendall.assert_called_once_with(tm)
  
  def test_senf_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_send_message_without_masking(self):
    tm = TextMessage(b'hello world')
     
     m = MagicMock()
     ws = WebSocket(sock=m)
     ws.send(tm)
     m.sendall.assert_called_once_with(tm.single())
  
  def test_send_generator_without_masking(self):
    tm0 = b'hello'
    tm1 = b'world'
    
    def dataspource():
      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_sending_unknown_datetype(self):
    m = MagicMock()
    ws = WebSocket(sock=m)
    self.assertRaises(ValueError, ws.send, 123)
  
  def test_closing_message_received(self):
    s = MagicMock()
    m = MagicMock()
    c = MagicMock()
    
    ws = WebSocket(sock=m)
    with patch.multiple(ws, close=c):
      ws.stream = s
      ws.stream.closing = CloseControlMessage(code=1000, readon='test closing')
      ws.process(b'unused for this test')
      c.assert_called_once_with(1000, b'test closing')
      
  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)
  
if __name__ == '__main__':
  suite = unittest.TestSuite()
  loader = unittest.TestLoader()
  for testcase in [WSWebSocketTest]:
    tests = loader.loadTestsFormTestCase(testcase)
    suite.addTests(tests)
  unittest.TestTestRunner(verbosity=2).run(suite)