diff --git a/lark_oapi/ws/client.py b/lark_oapi/ws/client.py index 8ee991838..39e2a1f9c 100644 --- a/lark_oapi/ws/client.py +++ b/lark_oapi/ws/client.py @@ -28,6 +28,10 @@ from lark_oapi.ws.pb.google.protobuf.internal.containers import RepeatedCompositeFieldContainer from lark_oapi.ws.pb.pbbp2_pb2 import Frame + +_ENDPOINT_CONNECT_TIMEOUT = 10 +_ENDPOINT_READ_TIMEOUT = 30 + try: loop = asyncio.get_event_loop() except RuntimeError: @@ -262,6 +266,7 @@ def _get_conn_url(self) -> str: url, headers=headers, json=body, + timeout=(_ENDPOINT_CONNECT_TIMEOUT, _ENDPOINT_READ_TIMEOUT), ) if response.status_code != http.HTTPStatus.OK: msg = "system busy" diff --git a/lark_oapi/ws/tests/test_websockets_compat.py b/lark_oapi/ws/tests/test_websockets_compat.py index 538d1ee94..d5fc19f6f 100644 --- a/lark_oapi/ws/tests/test_websockets_compat.py +++ b/lark_oapi/ws/tests/test_websockets_compat.py @@ -52,10 +52,11 @@ def test_parse_ws_connection_exception_keeps_legacy_headers_behavior(): def test_get_conn_url_sends_custom_headers(monkeypatch): captured = {} - def fake_post(url, *, headers=None, json=None): + def fake_post(url, *, headers=None, json=None, timeout=None): captured["url"] = url captured["headers"] = headers captured["json"] = json + captured["timeout"] = timeout return SimpleNamespace( status_code=200, content=b'{"code":0,"data":{"URL":"ws://example.test/callback?device_id=device&service_id=42"}}', @@ -76,6 +77,7 @@ def fake_post(url, *, headers=None, json=None): assert client._get_conn_url() == "ws://example.test/callback?device_id=device&service_id=42" assert captured["url"] == client._domain + ws_client.GEN_ENDPOINT_URI assert captured["json"] == {"AppID": "app_id", "AppSecret": "app_secret"} + assert captured["timeout"] == (10, 30) assert captured["headers"]["x-tt-env"] == "boe" assert captured["headers"]["x-use-ppe"] == "1" assert captured["headers"]["locale"] == "zh"