def test_constructor(self):
'''Make sure we can create a TCP object'''
self.comm1 = TCPCommunicator(8998)
self.assertNotEqual(self.comm1, None)
self.comm1.close()
self.comm1 = TCPCommunicator(8998)
self.assertNotEqual(self.comm1, None)
self.comm1.close()
def test_send_no_connect(self):
'''Send a message before connecting'''
self.comm1 = TCPCommunicator(7777)
self.comm1.listen()
self.comm1.send('localhost', b'ABC1234', b'99999')
time.sleep(0.1)
data = self.comm1.get('127.0.0.1', b'9999')
self.assertEqual(data, b'ABC1234')
def test_recv_get_tcp(self):
'''Ensure we can put a packet into the data store'''
self.comm1 = TCPCommunicator(8998)
tag = 'abcd'.encode('utf-8')[0:4]
data = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'.encode('utf-8')
pkt = tag + data
self.comm1.receive_tcp(pkt, '127.0.0.1')
self.assertEqual(self.comm1.get('127.0.0.1', tag), data)
self.comm1.close()
def test_listen(self):
'''Make sure the listening thread gets created'''
self.comm1 = TCPCommunicator(8998)
self.comm1.listen()
self.assertEqual(self.comm1.is_listening, True)
self.assertNotEqual(self.comm1.listen_thread, None)
self.comm1.close()
self.assertNotEqual(self.comm1.is_listening, True)
self.assertEqual(self.comm1.listen_thread, None)
def test_multiple_connections(self):
'''Try connecting to the same host multiple times'''
self.comm1 = TCPCommunicator(9090)
self.comm1.listen()
self.comm1.connect('localhost')
self.comm1.send('localhost', b'12345', b'1234')
time.sleep(0.2)
data = self.comm1.get('127.0.0.1', b'1234')
self.assertEqual(data, b'12345')
self.comm1.close()
def test_run_tcp(self, mock_conn, mock_sock):
'''Make sure that we create all of our threads'''
self.comm1 = TCPCommunicator(8998)
self.comm1.is_listening = True
self.comm1.listen_sock = socket.socket(socket.AF_INET,
socket.SOCK_STREAM)
thd = threading.Thread(target=self.comm1._run_tcp,
args=(self.comm1.listen_sock, '0.0.0.0',
self.comm1.port))
thd.start()
time.sleep(0.1)
self.comm1.is_listening = False
thd.join()
self.assertTrue(mock_sock.called)
self.assertTrue(mock_conn.called)
self.assertTrue(
isinstance(self.comm1.connections['127.0.0.1'], MagicMock))
self.comm1.close()
def test_run_connect(self, mock_recv):
'''Make sure that we can receive data'''
self.comm1 = TCPCommunicator(8998)
self.comm1.connections['127.0.0.1'] = socket.socket(
socket.AF_INET, socket.SOCK_STREAM)
self.comm1.is_listening = True
self.comm1.receive_tcp = MagicMock()
self.comm1.receive_tcp.return_value = -1
thd = threading.Thread(target=self.comm1._run_connect,
args=(self.comm1.connections['127.0.0.1'],
'127.0.0.1'))
thd.start()
time.sleep(0.1)
self.comm1.is_listening = False
thd.join()
self.assertTrue(self.comm1.receive_tcp.called)
self.assertTrue(mock_recv.called)
self.assertTrue('127.0.0.1' not in self.comm1.connections)
self.comm1.close()
def setUp(self):
if self.comm1 is None:
self.comm1 = TCPCommunicator(9090)
else:
self.comm1.close()
self.comm1 = TCPCommunicator(9090)
class TCPCommTest(unittest.TestCase):
def __init__(self, *args, **kwargs):
self.comm1 = None
super().__init__(*args, **kwargs)
def setUp(self):
if self.comm1 is None:
self.comm1 = TCPCommunicator(9090)
else:
self.comm1.close()
self.comm1 = TCPCommunicator(9090)
def tearDown(self):
if self.comm1 is not None:
self.comm1.close()
def test_multiple_connections(self):
'''Try connecting to the same host multiple times'''
self.comm1 = TCPCommunicator(9090)
self.comm1.listen()
self.comm1.connect('localhost')
self.comm1.send('localhost', b'12345', b'1234')
time.sleep(0.2)
data = self.comm1.get('127.0.0.1', b'1234')
self.assertEqual(data, b'12345')
self.comm1.close()
def test_send_no_connect(self):
'''Send a message before connecting'''
self.comm1 = TCPCommunicator(7777)
self.comm1.listen()
self.comm1.send('localhost', b'ABC1234', b'99999')
time.sleep(0.1)
data = self.comm1.get('127.0.0.1', b'9999')
self.assertEqual(data, b'ABC1234')
def test_constructor(self):
'''Make sure we can create a TCP object'''
self.comm1 = TCPCommunicator(8998)
self.assertNotEqual(self.comm1, None)
self.comm1.close()
self.comm1 = TCPCommunicator(8998)
self.assertNotEqual(self.comm1, None)
self.comm1.close()
def test_listen(self):
'''Make sure the listening thread gets created'''
self.comm1 = TCPCommunicator(8998)
self.comm1.listen()
self.assertEqual(self.comm1.is_listening, True)
self.assertNotEqual(self.comm1.listen_thread, None)
self.comm1.close()
self.assertNotEqual(self.comm1.is_listening, True)
self.assertEqual(self.comm1.listen_thread, None)
@patch('socket.socket.accept',
return_value=(MagicMock(), ('127.0.0.1', 1234)))
@patch('pymessage.TCPCommunicator._run_connect', return_value=-1)
def test_run_tcp(self, mock_conn, mock_sock):
'''Make sure that we create all of our threads'''
self.comm1 = TCPCommunicator(8998)
self.comm1.is_listening = True
self.comm1.listen_sock = socket.socket(socket.AF_INET,
socket.SOCK_STREAM)
thd = threading.Thread(target=self.comm1._run_tcp,
args=(self.comm1.listen_sock, '0.0.0.0',
self.comm1.port))
thd.start()
time.sleep(0.1)
self.comm1.is_listening = False
thd.join()
self.assertTrue(mock_sock.called)
self.assertTrue(mock_conn.called)
self.assertTrue(
isinstance(self.comm1.connections['127.0.0.1'], MagicMock))
self.comm1.close()
@patch(
'socket.socket.recv',
side_effect=[struct.pack('!I', 5), 'hello world'.encode('utf-8'), b''])
def test_run_connect(self, mock_recv):
'''Make sure that we can receive data'''
self.comm1 = TCPCommunicator(8998)
self.comm1.connections['127.0.0.1'] = socket.socket(
socket.AF_INET, socket.SOCK_STREAM)
self.comm1.is_listening = True
self.comm1.receive_tcp = MagicMock()
self.comm1.receive_tcp.return_value = -1
thd = threading.Thread(target=self.comm1._run_connect,
args=(self.comm1.connections['127.0.0.1'],
'127.0.0.1'))
thd.start()
time.sleep(0.1)
self.comm1.is_listening = False
thd.join()
self.assertTrue(self.comm1.receive_tcp.called)
self.assertTrue(mock_recv.called)
self.assertTrue('127.0.0.1' not in self.comm1.connections)
self.comm1.close()
def test_recv_get_tcp(self):
'''Ensure we can put a packet into the data store'''
self.comm1 = TCPCommunicator(8998)
tag = 'abcd'.encode('utf-8')[0:4]
data = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'.encode('utf-8')
pkt = tag + data
self.comm1.receive_tcp(pkt, '127.0.0.1')
self.assertEqual(self.comm1.get('127.0.0.1', tag), data)
self.comm1.close()
def test_recv_n_bytes(self):
'''Make sure that the function to receive at most "n" bytes works correctly
Used for message delimiting'''
with patch('socket.socket.recv', side_effect=[b'1', b'2', b'3', b'4']):
_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
msg = comm.recv_n_bytes(_sock, 4)
self.assertEqual(msg, b'1234')
_sock.close()
def recv(n):
if n < 8:
return b'tcp_sock'[:n]
else:
return b'tcp_sock'
with patch('socket.socket.recv', side_effect=recv):
_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
msg = comm.recv_n_bytes(_sock, 5)
self.assertEqual(msg, b'tcp_s')
msg = comm.recv_n_bytes(_sock, 25)
self.assertEqual(msg, (b'tcp_sock' * 5)[:25])
with self.assertRaises(ValueError):
msg = comm.recv_n_bytes(_sock, -1)
_sock.close()
with patch('socket.socket.recv', return_value=struct.pack('!I', 512)):
_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
msg = comm.recv_n_bytes(_sock, 4)
self.assertEqual(struct.unpack('!I', msg)[0], 512)
_sock.close()
with patch('socket.socket.recv',
side_effect=[b'\x00', b'\x00', b'\x02', b'\x00']):
_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
msg = comm.recv_n_bytes(_sock, 4)
self.assertEqual(struct.unpack('!I', msg)[0], 512)
_sock.close()