lib/py/test/test_socket.py

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
#

import errno
import time
import unittest

from test_sslsocket import ServerAcceptor

import _import_local_thrift  # noqa

from thrift.transport.TSocket import TServerSocket
from thrift.transport.TSocket import TSocket
from thrift.transport.TTransport import TTransportException


class TSocketTest(unittest.TestCase):
    def test_failed_connection_raises_exception(self):
        sock = TSocket(host="localhost", port=60606) # unused port
        with self.assertRaises(TTransportException) as ctx:
            sock.open()
        exc = ctx.exception
        self.assertEqual(exc.type, TTransportException.NOT_OPEN)
        self.assertIn("Could not connect to any of", exc.message)
        self.assertIsNotNone(exc.inner)
        self.assertIn("Connection refused", str(exc.inner))

    def test_socket_readtimeout_exception(self):
        acc = ServerAcceptor(TServerSocket(port=0))
        acc.start()

        sock = TSocket(host="localhost", port=acc.port)
        sock.open()
        sock.setTimeout(1)
        sock.write(b"sleep")

        with self.assertRaises(TTransportException) as ctx:
            sock.read(5)
        exc = ctx.exception
        self.assertEqual(exc.message, "read timeout")

        acc.client.close() # this also blocks until the other thread is done
        acc.close()
        sock.close()

    def test_isOpen_checks_for_readability(self):
        # https://docs.python.org/3/library/socket.html#notes-on-socket-timeouts
        # https://docs.python.org/3/library/socket.html#socket.socket.settimeout
        timeouts = [
            None,  # blocking mode
            0,  # non-blocking mode
            1.0,  # timeout mode
        ]

        for timeout in timeouts:
            acc = ServerAcceptor(TServerSocket(port=0))
            acc.start()

            sock = TSocket(host="localhost", port=acc.port)
            self.assertFalse(sock.isOpen())
            sock.open()
            sock.setTimeout(timeout)

            # the socket shows as open immediately after connecting
            self.assertTrue(sock.isOpen())

            # and remains open during usage
            sock.write(b"hello")
            self.assertTrue(sock.isOpen())
            while True:
                try:
                    sock.read(5)
                except TTransportException as exc:
                    if exc.inner.errno == errno.EAGAIN:
                        # try again when we're in non-blocking mode
                        continue
                    raise
                break
            self.assertTrue(sock.isOpen())

            # once the server side closes, it no longer shows open
            acc.client.close()  # this also blocks until the other thread is done
            acc.close()

            self.assertIsNotNone(sock.handle)
            # Give the kernel a moment to propagate FIN before asserting.
            deadline = time.monotonic() + 0.5
            while sock.isOpen() and time.monotonic() < deadline:
                time.sleep(0.01)
            self.assertFalse(sock.isOpen(), "socket still open after 0.5s")
            # after isOpen() returned False the socket should be closed (THRIFT-5813)
            self.assertIsNone(sock.handle)


if __name__ == "__main__":
    unittest.main()