class TestRecvBufferedSocket(unittest.TestCase):
def setUp(self):
self.chunk_size = 512
self.mock_socket = mock.Mock()
self.client_sock, self.server_sock = socket.socketpair()
self.buf_sock = RecvBufferedSocket(self.client_sock,
chunk_size=self.chunk_size)
self.mocked_buf_sock = RecvBufferedSocket(self.mock_socket,
chunk_size=self.chunk_size)
def tearDown(self):
self.buf_sock.close()
self.server_sock.close()
def test_getattr(self):
self.assertTrue(inspect.ismethod(self.buf_sock.recv))
self.assertFalse(inspect.isbuiltin(self.buf_sock.recv))
self.assertTrue(inspect.isbuiltin(self.buf_sock.connect))
def test_recv(self):
self.server_sock.sendall(b'A' * 300)
self.assertEqual(self.buf_sock.recv(1), b'A')
self.assertEqual(self.buf_sock.recv(200), b'A' * 200)
self.assertEqual(self.buf_sock.recv(99), b'A' * 99)
def test_recv_max_larger_than_buf(self):
double_chunk = self.chunk_size * 2
self.server_sock.sendall(b'A' * double_chunk)
self.assertEqual(self.buf_sock.recv(double_chunk), b'A' * double_chunk)
@mock.patch('selectors.DefaultSelector' if PY3 else 'select.select',
**PATCH_OPTS)
def test_recv_check_calls(self, mock_selector):
if PY3:
mock_selector = mock_selector.return_value.__enter__.return_value
mock_selector.register = mock.Mock()
# NB: the return value should actually be List[Tuple[SelectorKey, Events]], but our code only
# cares that _some_ event happened so we choose a simpler mock here. See
# https://docs.python.org/3/library/selectors.html#selectors.BaseSelector.select.
mock_selector.select = mock.Mock(return_value=[(1, "")])
else:
mock_selector.return_value = ([1], [], [])
self.mock_socket.recv.side_effect = [
b'A' * self.chunk_size, b'B' * self.chunk_size
]
self.assertEqual(self.mocked_buf_sock.recv(128), b'A' * 128)
self.mock_socket.recv.assert_called_once_with(self.chunk_size)
self.assertEqual(self.mocked_buf_sock.recv(128), b'A' * 128)
self.assertEqual(self.mocked_buf_sock.recv(128), b'A' * 128)
self.assertEqual(self.mocked_buf_sock.recv(128), b'A' * 128)
self.assertEqual(self.mock_socket.recv.call_count, 1)
self.assertEqual(self.mocked_buf_sock.recv(self.chunk_size),
b'B' * self.chunk_size)
self.assertEqual(self.mock_socket.recv.call_count, 2)
class TestRecvBufferedSocket(unittest.TestCase):
def setUp(self):
self.chunk_size = 512
self.mock_socket = mock.Mock()
self.client_sock, self.server_sock = socket.socketpair()
self.buf_sock = RecvBufferedSocket(self.client_sock,
chunk_size=self.chunk_size)
self.mocked_buf_sock = RecvBufferedSocket(self.mock_socket,
chunk_size=self.chunk_size)
def tearDown(self):
self.buf_sock.close()
self.server_sock.close()
def test_getattr(self):
self.assertTrue(inspect.ismethod(self.buf_sock.recv))
self.assertFalse(inspect.isbuiltin(self.buf_sock.recv))
self.assertTrue(inspect.isbuiltin(self.buf_sock.connect))
def test_recv(self):
self.server_sock.sendall(b'A' * 300)
self.assertEquals(self.buf_sock.recv(1), b'A')
self.assertEquals(self.buf_sock.recv(200), b'A' * 200)
self.assertEquals(self.buf_sock.recv(99), b'A' * 99)
def test_recv_max_larger_than_buf(self):
double_chunk = self.chunk_size * 2
self.server_sock.sendall(b'A' * double_chunk)
self.assertEquals(self.buf_sock.recv(double_chunk),
b'A' * double_chunk)
@mock.patch('pants.util.socket.select.select', **PATCH_OPTS)
def test_recv_check_calls(self, mock_select):
mock_select.return_value = ([1], [], [])
self.mock_socket.recv.side_effect = [
b'A' * self.chunk_size, b'B' * self.chunk_size
]
self.assertEquals(self.mocked_buf_sock.recv(128), b'A' * 128)
self.mock_socket.recv.assert_called_once_with(self.chunk_size)
self.assertEquals(self.mocked_buf_sock.recv(128), b'A' * 128)
self.assertEquals(self.mocked_buf_sock.recv(128), b'A' * 128)
self.assertEquals(self.mocked_buf_sock.recv(128), b'A' * 128)
self.assertEquals(self.mock_socket.recv.call_count, 1)
self.assertEquals(self.mocked_buf_sock.recv(self.chunk_size),
b'B' * self.chunk_size)
self.assertEquals(self.mock_socket.recv.call_count, 2)
def try_connect(self):
"""Creates a socket, connects it to the nailgun and returns the connected socket.
:returns: a connected `socket.socket`.
:raises: `NailgunClient.NailgunConnectionError` on failure to connect.
"""
sock = RecvBufferedSocket(socket.socket(socket.AF_INET, socket.SOCK_STREAM))
try:
sock.connect((self._host, self._port))
except (socket.error, socket.gaierror) as e:
logger.debug('Encountered socket exception {!r} when attempting connect to nailgun'.format(e))
sock.close()
raise self.NailgunConnectionError(
'Problem connecting to nailgun server at {}:{}: {!r}'.format(self._host, self._port, e))
else:
return sock
class TestRecvBufferedSocket(unittest.TestCase):
def setUp(self):
self.chunk_size = 512
self.mock_socket = mock.Mock()
self.client_sock, self.server_sock = socket.socketpair()
self.buf_sock = RecvBufferedSocket(self.client_sock, chunk_size=self.chunk_size)
self.mocked_buf_sock = RecvBufferedSocket(self.mock_socket, chunk_size=self.chunk_size)
def tearDown(self):
self.buf_sock.close()
self.server_sock.close()
def test_getattr(self):
self.assertTrue(inspect.ismethod(self.buf_sock.recv))
self.assertFalse(inspect.isbuiltin(self.buf_sock.recv))
self.assertTrue(inspect.isbuiltin(self.buf_sock.connect))
def test_recv(self):
self.server_sock.sendall(b'A' * 300)
self.assertEquals(self.buf_sock.recv(1), b'A')
self.assertEquals(self.buf_sock.recv(200), b'A' * 200)
self.assertEquals(self.buf_sock.recv(99), b'A' * 99)
def test_recv_max_larger_than_buf(self):
double_chunk = self.chunk_size * 2
self.server_sock.sendall(b'A' * double_chunk)
self.assertEquals(self.buf_sock.recv(double_chunk), b'A' * double_chunk)
@mock.patch('pants.util.socket.select.select', **PATCH_OPTS)
def test_recv_check_calls(self, mock_select):
mock_select.return_value = ([1], [], [])
self.mock_socket.recv.side_effect = [b'A' * self.chunk_size, b'B' * self.chunk_size]
self.assertEquals(self.mocked_buf_sock.recv(128), b'A' * 128)
self.mock_socket.recv.assert_called_once_with(self.chunk_size)
self.assertEquals(self.mocked_buf_sock.recv(128), b'A' * 128)
self.assertEquals(self.mocked_buf_sock.recv(128), b'A' * 128)
self.assertEquals(self.mocked_buf_sock.recv(128), b'A' * 128)
self.assertEquals(self.mock_socket.recv.call_count, 1)
self.assertEquals(self.mocked_buf_sock.recv(self.chunk_size), b'B' * self.chunk_size)
self.assertEquals(self.mock_socket.recv.call_count, 2)
def try_connect(self):
"""Creates a socket, connects it to the nailgun and returns the connected socket.
:returns: a connected `socket.socket`.
:raises: `NailgunClient.NailgunConnectionError` on failure to connect.
"""
sock = RecvBufferedSocket(socket.socket(socket.AF_INET, socket.SOCK_STREAM))
try:
sock.connect(self._address)
except (socket.error, socket.gaierror) as e:
logger.debug('Encountered socket exception {!r} when attempting connect to nailgun'.format(e))
sock.close()
raise self.NailgunConnectionError(
address=self._address_string,
pid=self.pid,
wrapped_exc=e,
traceback=sys.exc_info()[2]
)
else:
return sock
def try_connect(self):
"""Creates a socket, connects it to the nailgun and returns the connected socket.
:returns: a connected `socket.socket`.
:raises: `NailgunClient.NailgunConnectionError` on failure to connect.
"""
sock = RecvBufferedSocket(socket.socket(socket.AF_INET, socket.SOCK_STREAM))
try:
sock.connect(self._address)
except (socket.error, socket.gaierror) as e:
logger.debug('Encountered socket exception {!r} when attempting connect to nailgun'.format(e))
sock.close()
raise self.NailgunConnectionError(
address=self._address_string,
pid=self._maybe_last_pid(),
pgrp=self._maybe_last_pgrp(),
wrapped_exc=e,
)
else:
return sock
def try_connect(self):
"""Creates a socket, connects it to the nailgun and returns the connected socket.
:returns: a connected `socket.socket`.
:raises: `NailgunClient.NailgunConnectionError` on failure to connect.
"""
sock = RecvBufferedSocket(
sock=socket.socket(family=socket.AF_INET, type=socket.SOCK_STREAM))
try:
sock.connect(self._address)
except (socket.error, socket.gaierror) as e:
logger.debug(
"Encountered socket exception {!r} when attempting connect to nailgun"
.format(e))
sock.close()
raise self.NailgunConnectionError(
address=self._address_string,
pid=self._maybe_last_pid(),
pgrp=self._maybe_last_pgrp(),
wrapped_exc=e,
)
else:
return sock
