diff --git a/tests/test_websocket.py b/tests/test_websocket.py index 29cb0d2..856a9cc 100644 --- a/tests/test_websocket.py +++ b/tests/test_websocket.py @@ -65,7 +65,7 @@ class HeaderSockMock(SockMock): def __init__(self, fname): SockMock.__init__(self) path = os.path.join(os.path.dirname(__file__), fname) - self.add_packet(open(path).read()) + self.add_packet(open(path).read().encode('utf-8')) class WebSocketTest(unittest.TestCase): @@ -165,7 +165,7 @@ class WebSocketTest(unittest.TestCase): def testWSKey(self): key = ws._create_sec_websocket_key() self.assert_(key != 24) - self.assert_("¥n" not in key) + self.assert_(six.u("¥n") not in key) def testWsUtils(self): sock = ws.WebSocket() @@ -227,7 +227,8 @@ class WebSocketTest(unittest.TestCase): something = "\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc" s.add_packet(something) data = sock.recv() - self.assertEquals(data.decode('utf-8'), "こんにちは") + data = data.decode('utf-8') + self.assertEquals(data, u"こんにちは") s.add_packet("\x81\x85abcd)\x07\x0f\x08\x0e") data = sock.recv() diff --git a/websocket/__init__.py b/websocket/__init__.py index da4e87b..2e5e573 100644 --- a/websocket/__init__.py +++ b/websocket/__init__.py @@ -347,7 +347,11 @@ class ABNF(object): def _get_masked(self, mask_key): s = ABNF.mask(mask_key, self.data) - return six.b(mask_key) + s + + if isinstance(mask_key, six.text_type): + mask_key = mask_key.encode('utf-8') + + return mask_key + s @staticmethod def mask(mask_key, data): @@ -564,6 +568,7 @@ class WebSocket(object): while True: line = self._recv_line() + line = line.decode('utf-8') if line == "\r\n" or line == "\n": break line = line.strip() @@ -595,6 +600,7 @@ class WebSocket(object): opcode: operation code to send. Please see OPCODE_XXX. """ + frame = ABNF.create_frame(payload, opcode) return self.send_frame(frame) @@ -870,9 +876,9 @@ class WebSocket(object): while True: c = self._recv(1) line.append(c) - if c == "\n": + if c == six.b("\n"): break - return "".join(line) + return six.b("").join(line) class WebSocketApp(object): @@ -930,6 +936,7 @@ class WebSocketApp(object): data: message to send. If you set opcode to OPCODE_TEXT, data must be utf-8 string or unicode. opcode: operation code of data. default is OPCODE_TEXT. """ + if self.sock.send(data, opcode) == 0: raise WebSocketConnectionClosedException()