Restore removed tests and make Python 3 compatible

This restores tests removed in 23beb7d43e

1 files changed, 36 insertions, 2 deletions

diff --git a/tests/websocket_test.py b/tests/websocket_test.py
index 4b2aa4f..62d59a9 100644
--- a/tests/websocket_test.py
+++ b/tests/websocket_test.py
@@ -6,9 +6,9 @@ from eventlet import event
 from eventlet import greenio
 from eventlet.green import httplib
 from eventlet.support import six
-from eventlet.websocket import WebSocketWSGI
+from eventlet.websocket import WebSocket, WebSocketWSGI
 
-from tests import certificate_file, private_key_file
+from tests import certificate_file, LimitedTestCase, mock, private_key_file
 from tests import skip_if_no_ssl
 from tests.wsgi_test import _TestBase
 
@@ -566,3 +566,37 @@ class TestWebSocketSSL(_TestBase):
         greenio.shutdown_safe(sock)
         sock.close()
         eventlet.sleep(0.01)
+
+
+class TestWebSocketObject(LimitedTestCase):
+
+    def setUp(self):
+        self.mock_socket = s = mock.Mock()
+        self.environ = env = dict(HTTP_ORIGIN='http://localhost', HTTP_WEBSOCKET_PROTOCOL='ws',
+                                  PATH_INFO='test')
+
+        self.test_ws = WebSocket(s, env)
+        super(TestWebSocketObject, self).setUp()
+
+    def test_recieve(self):
+        ws = self.test_ws
+        ws.socket.recv.return_value = b'\x00hello\xFF'
+        self.assertEqual(ws.wait(), 'hello')
+        self.assertEqual(ws._buf, b'')
+        self.assertEqual(len(ws._msgs), 0)
+        ws.socket.recv.return_value = b''
+        self.assertEqual(ws.wait(), None)
+        self.assertEqual(ws._buf, b'')
+        self.assertEqual(len(ws._msgs), 0)
+
+    def test_send_to_ws(self):
+        ws = self.test_ws
+        ws.send(u'hello')
+        assert ws.socket.sendall.called_with("\x00hello\xFF")
+        ws.send(10)
+        assert ws.socket.sendall.called_with("\x0010\xFF")
+
+    def test_close_ws(self):
+        ws = self.test_ws
+        ws.close()
+        assert ws.socket.shutdown.called_with(True)