codekingpro/portable-devtools

1# -*- coding: utf-8 -*-2#3import os4import os.path5import socket6import unittest7from base64 import decodebytes as base64decode8 9import websocket as ws10from websocket._exceptions import (11    WebSocketBadStatusException,12    WebSocketAddressException,13    WebSocketException,14)15from websocket._handshake import _create_sec_websocket_key16from websocket._handshake import _validate as _validate_header17from websocket._http import read_headers18from websocket._utils import validate_utf819 20"""21test_websocket.py22websocket - WebSocket client library for Python23 24Copyright 2025 engn33r25 26Licensed under the Apache License, Version 2.0 (the "License");27you may not use this file except in compliance with the License.28You may obtain a copy of the License at29 30    http://www.apache.org/licenses/LICENSE-2.031 32Unless required by applicable law or agreed to in writing, software33distributed under the License is distributed on an "AS IS" BASIS,34WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.35See the License for the specific language governing permissions and36limitations under the License.37"""38 39try:40    import ssl41except ImportError:42    # dummy class of SSLError for ssl none-support environment.43    class SSLError(Exception):44        pass45 46 47# Skip test to access the internet unless TEST_WITH_INTERNET == 148TEST_WITH_INTERNET = os.environ.get("TEST_WITH_INTERNET", "0") == "1"49# Skip tests relying on local websockets server unless LOCAL_WS_SERVER_PORT != -150LOCAL_WS_SERVER_PORT = os.environ.get("LOCAL_WS_SERVER_PORT", "-1")51TEST_WITH_LOCAL_SERVER = LOCAL_WS_SERVER_PORT != "-1"52TRACEABLE = True53 54 55def create_mask_key(_):56    return "abcd"57 58 59class SockMock:60    def __init__(self):61        self.data = []62        self.sent = []63 64    def add_packet(self, data):65        self.data.append(data)66 67    def gettimeout(self):68        return None69 70    def recv(self, bufsize):71        if self.data:72            e = self.data.pop(0)73            if isinstance(e, Exception):74                raise e75            if len(e) > bufsize:76                self.data.insert(0, e[bufsize:])77            return e[:bufsize]78 79    def send(self, data):80        self.sent.append(data)81        return len(data)82 83    def close(self):84        pass85 86 87class HeaderSockMock(SockMock):88    def __init__(self, fname):89        SockMock.__init__(self)90        path = os.path.join(os.path.dirname(__file__), fname)91        with open(path, "rb") as f:92            self.add_packet(f.read())93 94 95class WebSocketTest(unittest.TestCase):96    def setUp(self):97        ws.enableTrace(TRACEABLE)98 99    def tearDown(self):100        pass101 102    def test_default_timeout(self):103        self.assertEqual(ws.getdefaulttimeout(), None)104        ws.setdefaulttimeout(10)105        self.assertEqual(ws.getdefaulttimeout(), 10)106        ws.setdefaulttimeout(None)107 108    def test_ws_key(self):109        key = _create_sec_websocket_key()110        self.assertTrue(key != 24)111        self.assertTrue("¥n" not in key)112 113    def test_nonce(self):114        """WebSocket key should be a random 16-byte nonce."""115        key = _create_sec_websocket_key()116        nonce = base64decode(key.encode("utf-8"))117        self.assertEqual(16, len(nonce))118 119    def test_ws_utils(self):120        key = "c6b8hTg4EeGb2gQMztV1/g=="121        required_header = {122            "upgrade": "websocket",123            "connection": "upgrade",124            "sec-websocket-accept": "Kxep+hNu9n51529fGidYu7a3wO0=",125        }126        self.assertEqual(_validate_header(required_header, key, None), (True, None))127 128        header = required_header.copy()129        header["upgrade"] = "http"130        self.assertEqual(_validate_header(header, key, None), (False, None))131        del header["upgrade"]132        self.assertEqual(_validate_header(header, key, None), (False, None))133 134        header = required_header.copy()135        header["connection"] = "something"136        self.assertEqual(_validate_header(header, key, None), (False, None))137        del header["connection"]138        self.assertEqual(_validate_header(header, key, None), (False, None))139 140        header = required_header.copy()141        header["sec-websocket-accept"] = "something"142        self.assertEqual(_validate_header(header, key, None), (False, None))143        del header["sec-websocket-accept"]144        self.assertEqual(_validate_header(header, key, None), (False, None))145 146        header = required_header.copy()147        header["sec-websocket-protocol"] = "sub1"148        self.assertEqual(149            _validate_header(header, key, ["sub1", "sub2"]), (True, "sub1")150        )151        # This case will print out a logging error using the error() function, but that is expected152        self.assertEqual(_validate_header(header, key, ["sub2", "sub3"]), (False, None))153 154        header = required_header.copy()155        header["sec-websocket-protocol"] = "sUb1"156        self.assertEqual(157            _validate_header(header, key, ["Sub1", "suB2"]), (True, "sub1")158        )159 160        header = required_header.copy()161        # This case will print out a logging error using the error() function, but that is expected162        self.assertEqual(_validate_header(header, key, ["Sub1", "suB2"]), (False, None))163 164    def test_read_header(self):165        status, header, _ = read_headers(HeaderSockMock("data/header01.txt"))166        self.assertEqual(status, 101)167        self.assertEqual(header["connection"], "Upgrade")168 169        status, header, _ = read_headers(HeaderSockMock("data/header03.txt"))170        self.assertEqual(status, 101)171        self.assertEqual(header["connection"], "Upgrade, Keep-Alive")172 173        HeaderSockMock("data/header02.txt")174        self.assertRaises(175            ws.WebSocketException, read_headers, HeaderSockMock("data/header02.txt")176        )177 178    def test_send(self):179        # TODO: add longer frame data180        sock = ws.WebSocket()181        sock.set_mask_key(create_mask_key)182        s = sock.sock = HeaderSockMock("data/header01.txt")183        sock.send("Hello")184        self.assertEqual(s.sent[0], b"\x81\x85abcd)\x07\x0f\x08\x0e")185 186        sock.send("こんにちは")187        self.assertEqual(188            s.sent[1],189            b"\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc",190        )191 192        #        sock.send("x" * 5000)193        #        self.assertEqual(s.sent[1], b'\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc")194 195        self.assertEqual(sock.send_binary(b"1111111111101"), 19)196 197    def test_recv(self):198        # TODO: add longer frame data199        sock = ws.WebSocket()200        s = sock.sock = SockMock()201        something = (202            b"\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc"203        )204        s.add_packet(something)205        data = sock.recv()206        self.assertEqual(data, "こんにちは")207 208        s.add_packet(b"\x81\x85abcd)\x07\x0f\x08\x0e")209        data = sock.recv()210        self.assertEqual(data, "Hello")211 212    @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")213    def test_iter(self):214        count = 2215        s = ws.create_connection("wss://api.bitfinex.com/ws/2")216        s.send('{"event": "subscribe", "channel": "ticker"}')217        for _ in s:218            count -= 1219            if count == 0:220                break221 222    @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")223    def test_next(self):224        sock = ws.create_connection("wss://api.bitfinex.com/ws/2")225        self.assertEqual(str, type(next(sock)))226 227    def test_internal_recv_strict(self):228        sock = ws.WebSocket()229        s = sock.sock = SockMock()230        s.add_packet(b"foo")231        s.add_packet(socket.timeout())232        s.add_packet(b"bar")233        # s.add_packet(SSLError("The read operation timed out"))234        s.add_packet(b"baz")235        with self.assertRaises(ws.WebSocketTimeoutException):236            sock.frame_buffer.recv_strict(9)237        #     with self.assertRaises(SSLError):238        #         data = sock._recv_strict(9)239        data = sock.frame_buffer.recv_strict(9)240        self.assertEqual(data, b"foobarbaz")241        with self.assertRaises(ws.WebSocketConnectionClosedException):242            sock.frame_buffer.recv_strict(1)243 244    def test_recv_timeout(self):245        sock = ws.WebSocket()246        s = sock.sock = SockMock()247        s.add_packet(b"\x81")248        s.add_packet(socket.timeout())249        s.add_packet(b"\x8dabcd\x29\x07\x0f\x08\x0e")250        s.add_packet(socket.timeout())251        s.add_packet(b"\x4e\x43\x33\x0e\x10\x0f\x00\x40")252        with self.assertRaises(ws.WebSocketTimeoutException):253            sock.recv()254        with self.assertRaises(ws.WebSocketTimeoutException):255            sock.recv()256        data = sock.recv()257        self.assertEqual(data, "Hello, World!")258        with self.assertRaises(ws.WebSocketConnectionClosedException):259            sock.recv()260 261    def test_recv_with_simple_fragmentation(self):262        sock = ws.WebSocket()263        s = sock.sock = SockMock()264        # OPCODE=TEXT, FIN=0, MSG="Brevity is "265        s.add_packet(b"\x01\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")266        # OPCODE=CONT, FIN=1, MSG="the soul of wit"267        s.add_packet(b"\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")268        data = sock.recv()269        self.assertEqual(data, "Brevity is the soul of wit")270        with self.assertRaises(ws.WebSocketConnectionClosedException):271            sock.recv()272 273    def test_recv_with_fire_event_of_fragmentation(self):274        sock = ws.WebSocket(fire_cont_frame=True)275        s = sock.sock = SockMock()276        # OPCODE=TEXT, FIN=0, MSG="Brevity is "277        s.add_packet(b"\x01\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")278        # OPCODE=CONT, FIN=0, MSG="Brevity is "279        s.add_packet(b"\x00\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")280        # OPCODE=CONT, FIN=1, MSG="the soul of wit"281        s.add_packet(b"\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")282 283        _, data = sock.recv_data()284        self.assertEqual(data, b"Brevity is ")285        _, data = sock.recv_data()286        self.assertEqual(data, b"Brevity is ")287        _, data = sock.recv_data()288        self.assertEqual(data, b"the soul of wit")289 290        # OPCODE=CONT, FIN=0, MSG="Brevity is "291        s.add_packet(b"\x80\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")292 293        with self.assertRaises(ws.WebSocketException):294            sock.recv_data()295 296        with self.assertRaises(ws.WebSocketConnectionClosedException):297            sock.recv()298 299    def test_close(self):300        sock = ws.WebSocket()301        sock.connected = True302        sock.close()303 304        sock = ws.WebSocket()305        s = sock.sock = SockMock()306        sock.connected = True307        s.add_packet(b"\x88\x80\x17\x98p\x84")308        sock.recv()309        self.assertEqual(sock.connected, False)310 311    def test_recv_cont_fragmentation(self):312        sock = ws.WebSocket()313        s = sock.sock = SockMock()314        # OPCODE=CONT, FIN=1, MSG="the soul of wit"315        s.add_packet(b"\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")316        self.assertRaises(ws.WebSocketException, sock.recv)317 318    def test_recv_with_prolonged_fragmentation(self):319        sock = ws.WebSocket()320        s = sock.sock = SockMock()321        # OPCODE=TEXT, FIN=0, MSG="Once more unto the breach, "322        s.add_packet(323            b"\x01\x9babcd.\x0c\x00\x01A\x0f\x0c\x16\x04B\x16\n\x15\rC\x10\t\x07C\x06\x13\x07\x02\x07\tNC"324        )325        # OPCODE=CONT, FIN=0, MSG="dear friends, "326        s.add_packet(b"\x00\x8eabcd\x05\x07\x02\x16A\x04\x11\r\x04\x0c\x07\x17MB")327        # OPCODE=CONT, FIN=1, MSG="once more"328        s.add_packet(b"\x80\x89abcd\x0e\x0c\x00\x01A\x0f\x0c\x16\x04")329        data = sock.recv()330        self.assertEqual(data, "Once more unto the breach, dear friends, once more")331        with self.assertRaises(ws.WebSocketConnectionClosedException):332            sock.recv()333 334    def test_recv_with_fragmentation_and_control_frame(self):335        sock = ws.WebSocket()336        sock.set_mask_key(create_mask_key)337        s = sock.sock = SockMock()338        # OPCODE=TEXT, FIN=0, MSG="Too much "339        s.add_packet(b"\x01\x89abcd5\r\x0cD\x0c\x17\x00\x0cA")340        # OPCODE=PING, FIN=1, MSG="Please PONG this"341        s.add_packet(b"\x89\x90abcd1\x0e\x06\x05\x12\x07C4.,$D\x15\n\n\x17")342        # OPCODE=CONT, FIN=1, MSG="of a good thing"343        s.add_packet(b"\x80\x8fabcd\x0e\x04C\x05A\x05\x0c\x0b\x05B\x17\x0c\x08\x0c\x04")344        data = sock.recv()345        self.assertEqual(data, "Too much of a good thing")346        with self.assertRaises(ws.WebSocketConnectionClosedException):347            sock.recv()348        self.assertEqual(349            s.sent[0], b"\x8a\x90abcd1\x0e\x06\x05\x12\x07C4.,$D\x15\n\n\x17"350        )351 352    @unittest.skipUnless(353        TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"354    )355    def test_websocket(self):356        s = ws.create_connection(f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}")357        self.assertNotEqual(s, None)358        s.send("Hello, World")359        result = s.next()360        s.fileno()361        self.assertEqual(result, "Hello, World")362 363        s.send("こにゃにゃちは、世界")364        result = s.recv()365        self.assertEqual(result, "こにゃにゃちは、世界")366        self.assertRaises(ValueError, s.send_close, -1, "")367        s.close()368 369    @unittest.skipUnless(370        TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"371    )372    def test_ping_pong(self):373        s = ws.create_connection(f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}")374        self.assertNotEqual(s, None)375        s.ping("Hello")376        s.pong("Hi")377        s.close()378 379    @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")380    def test_support_redirect(self):381        s = ws.WebSocket()382        self.assertRaises(WebSocketBadStatusException, s.connect, "ws://google.com/")383        # Need to find a URL that has a redirect code leading to a websocket384 385    @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")386    def test_secure_websocket(self):387        s = ws.create_connection("wss://api.bitfinex.com/ws/2")388        self.assertNotEqual(s, None)389        self.assertTrue(isinstance(s.sock, ssl.SSLSocket))390        self.assertEqual(s.getstatus(), 101)391        self.assertNotEqual(s.getheaders(), None)392        s.settimeout(10)393        self.assertEqual(s.gettimeout(), 10)394        self.assertEqual(s.getsubprotocol(), None)395        s.abort()396 397    @unittest.skipUnless(398        TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"399    )400    def test_websocket_with_custom_header(self):401        s = ws.create_connection(402            f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}",403            headers={"User-Agent": "PythonWebsocketClient"},404        )405        self.assertNotEqual(s, None)406        self.assertEqual(s.getsubprotocol(), None)407        s.send("Hello, World")408        result = s.recv()409        self.assertEqual(result, "Hello, World")410        self.assertRaises(ValueError, s.close, -1, "")411        s.close()412 413    @unittest.skipUnless(414        TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"415    )416    def test_after_close(self):417        s = ws.create_connection(f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}")418        self.assertNotEqual(s, None)419        s.close()420        self.assertRaises(ws.WebSocketConnectionClosedException, s.send, "Hello")421        self.assertRaises(ws.WebSocketConnectionClosedException, s.recv)422 423 424class SockOptTest(unittest.TestCase):425    @unittest.skipUnless(426        TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"427    )428    def test_sockopt(self):429        sockopt = ((socket.IPPROTO_TCP, socket.TCP_NODELAY, 1),)430        s = ws.create_connection(431            f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}", sockopt=sockopt432        )433        self.assertNotEqual(434            s.sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY), 0435        )436        s.close()437 438 439class UtilsTest(unittest.TestCase):440    def test_utf8_validator(self):441        state = validate_utf8(b"\xf0\x90\x80\x80")442        self.assertEqual(state, True)443        state = validate_utf8(444            b"\xce\xba\xe1\xbd\xb9\xcf\x83\xce\xbc\xce\xb5\xed\xa0\x80edited"445        )446        self.assertEqual(state, False)447        state = validate_utf8(b"")448        self.assertEqual(state, True)449 450 451class HandshakeTest(unittest.TestCase):452    @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")453    def test_http_ssl(self):454        websock1 = ws.WebSocket(455            sslopt={"cert_chain": ssl.get_default_verify_paths().capath},456            enable_multithread=False,457        )458        self.assertRaises(ValueError, websock1.connect, "wss://api.bitfinex.com/ws/2")459        websock2 = ws.WebSocket(sslopt={"certfile": "myNonexistentCertFile"})460        self.assertRaises(461            WebSocketException, websock2.connect, "wss://api.bitfinex.com/ws/2"462        )463 464    @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")465    def test_manual_headers(self):466        websock3 = ws.WebSocket(467            sslopt={468                "ca_certs": ssl.get_default_verify_paths().cafile,469                "ca_cert_path": ssl.get_default_verify_paths().capath,470            }471        )472        self.assertRaises(473            WebSocketBadStatusException,474            websock3.connect,475            "wss://api.bitfinex.com/ws/2",476            cookie="chocolate",477            origin="testing_websockets.com",478            host="echo.websocket.events/websocket-client-test",479            subprotocols=["testproto"],480            connection="Upgrade",481            header={482                "CustomHeader1": "123",483                "Cookie": "TestValue",484                "Sec-WebSocket-Key": "k9kFAUWNAMmf5OEMfTlOEA==",485                "Sec-WebSocket-Protocol": "newprotocol",486            },487        )488 489    def test_ipv6(self):490        websock2 = ws.WebSocket()491        self.assertRaises(ValueError, websock2.connect, "2001:4860:4860::8888")492 493    def test_bad_urls(self):494        websock3 = ws.WebSocket()495        self.assertRaises(ValueError, websock3.connect, "ws//example.com")496        self.assertRaises(WebSocketAddressException, websock3.connect, "ws://example")497        self.assertRaises(ValueError, websock3.connect, "example.com")498 499 500if __name__ == "__main__":501    unittest.main()502