diff --git a/tests/test_socket.py b/tests/test_socket.py index 7f0c9c6..39f9949 100644 --- a/tests/test_socket.py +++ b/tests/test_socket.py @@ -117,3 +117,27 @@ def test_server_socket_close(): conn.close() server_socket.close() + + +def test_client_socket_set_timeout(): + server_socket = TServerSocket(host="localhost", port=12345, + client_timeout=100) + server_socket.listen() + + client_socket = TSocket(host="localhost", port=12345, socket_timeout=100) + client_socket.open() + + conn = server_socket.accept() + + assert client_socket.sock.gettimeout() == 100 / 1000 + assert conn.sock.gettimeout() == 100 / 1000 + assert conn.sock.gettimeout() == 100 / 1000 + + client_socket.set_timeout(200) + conn.set_timeout(200) + assert client_socket.sock.gettimeout() == 200 / 1000 + assert conn.sock.gettimeout() == 200 / 1000 + + conn.close() + client_socket.close() + server_socket.close() diff --git a/thriftpy/transport/socket.py b/thriftpy/transport/socket.py index 7958425..8e7b068 100644 --- a/thriftpy/transport/socket.py +++ b/thriftpy/transport/socket.py @@ -78,6 +78,9 @@ def set_timeout(self, ms): self.socket_timeout = ms / 1000 if (ms and ms > 0) else None self.connect_timeout = self.socket_timeout + if self.sock is not None: + self.sock.settimeout(self.socket_timeout) + def is_open(self): return bool(self.sock)