def test_init_and_reg_shutdown(self):
"""
Test initialization of a socket manager.
"""
self.assertFalse(self.fake_socket.connected)
# Callbacks should never trigger during proper setup and shutdown
nop_called = False
def nop(*args):
nonlocal nop_called # noqa 999 we don't support py2
nop_called = True
socket_manager = SocketManager(
'https://127.0.0.1',
self.fake_socket.port,
nop,
nop,
nop,
TASK_GROUP_ID_1,
0.3,
nop,
)
self.assertTrue(self.fake_socket.connected)
self.assertFalse(nop_called)
# Test shutdown
self.assertFalse(self.fake_socket.disconnected)
self.assertFalse(socket_manager.is_shutdown)
self.assertTrue(socket_manager.alive)
socket_manager.shutdown()
time.sleep(0.3)
self.assertTrue(self.fake_socket.disconnected)
self.assertTrue(socket_manager.is_shutdown)
self.assertFalse(nop_called)
def setUp(self):
self.fake_socket = MockSocket()
time.sleep(0.3)
self.agent1 = MockAgent(
TEST_HIT_ID_1, TEST_ASSIGNMENT_ID_1, TEST_WORKER_ID_1, TASK_GROUP_ID_1
)
self.agent2 = MockAgent(
TEST_HIT_ID_2, TEST_ASSIGNMENT_ID_2, TEST_WORKER_ID_2, TASK_GROUP_ID_1
)
self.alive_packet = None
self.message_packet = None
self.dead_worker_id = None
self.dead_assignment_id = None
self.server_died = False
self.socket_manager = SocketManager(
'https://127.0.0.1',
3030,
self.on_alive,
self.on_message,
self.on_worker_death,
TASK_GROUP_ID_1,
1,
self.on_server_death,
)
def setUp(self):
self.AGENT_ALIVE_PACKET = Packet(
MESSAGE_ID_1,
data_model.AGENT_ALIVE,
self.SENDER_ID,
self.WORLD_ID,
self.ASSIGNMENT_ID,
self.DATA,
self.CONVERSATION_ID,
)
self.MESSAGE_SEND_PACKET_1 = Packet(
MESSAGE_ID_2,
data_model.WORLD_MESSAGE,
self.WORLD_ID,
self.SENDER_ID,
self.ASSIGNMENT_ID,
self.DATA,
self.CONVERSATION_ID,
)
self.MESSAGE_SEND_PACKET_2 = Packet(
MESSAGE_ID_3,
data_model.MESSAGE_BATCH,
self.WORLD_ID,
self.SENDER_ID,
self.ASSIGNMENT_ID,
self.DATA,
self.CONVERSATION_ID,
)
self.MESSAGE_SEND_PACKET_3 = Packet(
MESSAGE_ID_4,
data_model.MESSAGE_BATCH,
self.WORLD_ID,
self.SENDER_ID,
self.ASSIGNMENT_ID,
self.DATA,
self.CONVERSATION_ID,
)
self.fake_socket = MockSocket()
time.sleep(0.3)
self.alive_packet = None
self.message_packet = None
self.dead_worker_id = None
self.dead_assignment_id = None
self.server_died = False
self.socket_manager = SocketManager(
'https://127.0.0.1',
self.fake_socket.port,
self.on_alive,
self.on_message,
self.on_worker_death,
TASK_GROUP_ID_1,
1,
self.on_server_death,
)
def test_init_world_dead(self):
"""
Test initialization of a socket manager with a failed startup.
"""
self.assertFalse(self.fake_socket.connected)
self.fake_socket.close()
# Callbacks should never trigger during proper setup and shutdown
nop_called = False
def nop(*args):
nonlocal nop_called # noqa 999 we don't support py2
nop_called = True
server_death_called = False
def server_death(*args):
nonlocal server_death_called
server_death_called = True
with self.assertRaises(ConnectionRefusedError):
socket_manager = SocketManager(
'https://127.0.0.1',
self.fake_socket.port,
nop,
nop,
nop,
TASK_GROUP_ID_1,
0.4,
server_death,
)
self.assertIsNone(socket_manager)
self.assertFalse(nop_called)
self.assertTrue(server_death_called)
def test_init_and_socket_shutdown_then_restart(self):
"""
Test restoring connection to a socket.
"""
self.assertFalse(self.fake_socket.connected)
# Callbacks should never trigger during proper setup and shutdown
nop_called = False
def nop(*args):
nonlocal nop_called # noqa 999 we don't support py2
nop_called = True
server_death_called = False
def server_death(*args):
nonlocal server_death_called
server_death_called = True
socket_manager = SocketManager(
'https://127.0.0.1',
self.fake_socket.port,
nop,
nop,
nop,
TASK_GROUP_ID_1,
0.4,
server_death,
)
self.assertTrue(self.fake_socket.connected)
self.assertFalse(nop_called)
self.assertFalse(server_death_called)
# Test shutdown
self.assertFalse(self.fake_socket.disconnected)
self.assertFalse(socket_manager.is_shutdown)
self.assertTrue(socket_manager.alive)
self.fake_socket.close()
self.assertEqualBy(lambda: socket_manager.alive, False, 8)
self.assertFalse(socket_manager.alive)
self.fake_socket = MockSocket()
self.assertEqualBy(lambda: socket_manager.alive, True, 4)
self.assertFalse(nop_called)
self.assertFalse(server_death_called)
socket_manager.shutdown()
def test_init_and_socket_shutdown(self):
'''Test initialization of a socket manager with a failed shutdown'''
self.assertFalse(self.fake_socket.connected)
# Callbacks should never trigger during proper setup and shutdown
nop_called = False
def nop(*args):
nonlocal nop_called # noqa 999 we don't support py2
nop_called = True
server_death_called = False
def server_death(*args):
nonlocal server_death_called
server_death_called = True
socket_manager = SocketManager(
'https://127.0.0.1',
self.fake_socket.port,
nop,
nop,
nop,
TASK_GROUP_ID_1,
0.4,
server_death,
)
self.assertTrue(self.fake_socket.connected)
self.assertFalse(nop_called)
self.assertFalse(server_death_called)
# Test shutdown
self.assertFalse(self.fake_socket.disconnected)
self.assertFalse(socket_manager.is_shutdown)
self.assertTrue(socket_manager.alive)
self.fake_socket.close()
self.assertEqualBy(lambda: socket_manager.alive, False, 8)
self.assertEqualBy(lambda: server_death_called, True, 20)
self.assertFalse(nop_called)
socket_manager.shutdown()
class TestSocketManagerMessageHandling(unittest.TestCase):
"""
Test sending messages to the world and then to each of two agents, along with
failure cases for each.
"""
def on_alive(self, packet):
self.alive_packet = packet
self.socket_manager.open_channel(packet.sender_id,
packet.assignment_id)
def on_message(self, packet):
self.message_packet = packet
def on_worker_death(self, worker_id, assignment_id):
self.dead_worker_id = worker_id
self.dead_assignment_id = assignment_id
def on_server_death(self):
self.server_died = True
def assertEqualBy(self, val_func, val, max_time):
start_time = time.time()
while val_func() != val:
assert (time.time() - start_time <
max_time), "Value was not attained in specified time"
time.sleep(0.1)
def setUp(self):
self.fake_socket = MockSocket()
time.sleep(0.3)
self.agent1 = MockAgent(TEST_HIT_ID_1, TEST_ASSIGNMENT_ID_1,
TEST_WORKER_ID_1, TASK_GROUP_ID_1)
self.agent2 = MockAgent(TEST_HIT_ID_2, TEST_ASSIGNMENT_ID_2,
TEST_WORKER_ID_2, TASK_GROUP_ID_1)
self.alive_packet = None
self.message_packet = None
self.dead_worker_id = None
self.dead_assignment_id = None
self.server_died = False
self.socket_manager = SocketManager(
'https://127.0.0.1',
3030,
self.on_alive,
self.on_message,
self.on_worker_death,
TASK_GROUP_ID_1,
1,
self.on_server_death,
)
def tearDown(self):
self.socket_manager.shutdown()
self.fake_socket.close()
def test_alive_send_and_disconnect(self):
message_packet = None
def on_msg(*args):
nonlocal message_packet
message_packet = args[0]
self.agent1.register_to_socket(self.fake_socket, on_msg)
self.assertIsNone(message_packet)
# Assert alive is registered
alive_id = self.agent1.send_alive()
self.assertIsNone(message_packet)
self.assertIsNone(self.message_packet)
self.assertEqualBy(lambda: self.alive_packet is None, False, 8)
self.assertEqual(self.alive_packet.id, alive_id)
# Test message send from agent
test_message_text_1 = 'test_message_text_1'
msg_id = self.agent1.send_message(test_message_text_1)
self.assertEqualBy(lambda: self.message_packet is None, False, 8)
self.assertEqual(self.message_packet.id, msg_id)
self.assertEqual(self.message_packet.data['text'], test_message_text_1)
# Test message send to agent
manager_message_id = 'message_id_from_manager'
test_message_text_2 = 'test_message_text_2'
message_send_packet = Packet(
manager_message_id,
data_model.MESSAGE_BATCH,
self.socket_manager.get_my_sender_id(),
TEST_WORKER_ID_1,
TEST_ASSIGNMENT_ID_1,
test_message_text_2,
't2',
)
self.socket_manager.queue_packet(message_send_packet)
self.assertEqualBy(lambda: message_packet is None, False, 8)
self.assertEqual(message_packet.id, manager_message_id)
self.assertEqual(message_packet.data, test_message_text_2)
self.assertIn(manager_message_id, self.socket_manager.packet_map)
# Test agent disconnect
self.agent1.send_disconnect()
self.assertEqualBy(lambda: self.dead_worker_id, TEST_WORKER_ID_1, 8)
self.assertEqual(self.dead_assignment_id, TEST_ASSIGNMENT_ID_1)
def test_one_agent_disconnect_other_alive(self):
message_packet = None
def on_msg(*args):
nonlocal message_packet
message_packet = args[0]
self.agent1.register_to_socket(self.fake_socket, on_msg)
self.agent2.register_to_socket(self.fake_socket, on_msg)
self.assertIsNone(message_packet)
# Assert alive is registered
self.agent1.send_alive()
self.agent2.send_alive()
self.assertIsNone(message_packet)
# Kill second agent
self.agent2.send_disconnect()
self.assertEqualBy(lambda: self.dead_worker_id, TEST_WORKER_ID_2, 8)
self.assertEqual(self.dead_assignment_id, TEST_ASSIGNMENT_ID_2)
# Run rest of tests
# Test message send from agent
test_message_text_1 = 'test_message_text_1'
msg_id = self.agent1.send_message(test_message_text_1)
self.assertEqualBy(lambda: self.message_packet is None, False, 8)
self.assertEqual(self.message_packet.id, msg_id)
self.assertEqual(self.message_packet.data['text'], test_message_text_1)
# Test message send to agent
manager_message_id = 'message_id_from_manager'
test_message_text_2 = 'test_message_text_2'
message_send_packet = Packet(
manager_message_id,
data_model.WORLD_MESSAGE,
self.socket_manager.get_my_sender_id(),
TEST_WORKER_ID_1,
TEST_ASSIGNMENT_ID_1,
test_message_text_2,
't2',
)
self.socket_manager.queue_packet(message_send_packet)
self.assertEqualBy(lambda: message_packet is None, False, 8)
self.assertEqual(message_packet.id, manager_message_id)
self.assertEqual(message_packet.data, test_message_text_2)
self.assertIn(manager_message_id, self.socket_manager.packet_map)
# Test agent disconnect
self.agent1.send_disconnect()
self.assertEqualBy(lambda: self.dead_worker_id, TEST_WORKER_ID_1, 8)
self.assertEqual(self.dead_assignment_id, TEST_ASSIGNMENT_ID_1)
class TestSocketManagerRoutingFunctionality(unittest.TestCase):
ID = 'ID'
SENDER_ID = 'SENDER_ID'
ASSIGNMENT_ID = 'ASSIGNMENT_ID'
DATA = 'DATA'
CONVERSATION_ID = 'CONVERSATION_ID'
ACK_FUNCTION = 'ACK_FUNCTION'
WORLD_ID = '[World_{}]'.format(TASK_GROUP_ID_1)
def on_alive(self, packet):
self.alive_packet = packet
def on_message(self, packet):
self.message_packet = packet
def on_worker_death(self, worker_id, assignment_id):
self.dead_worker_id = worker_id
self.dead_assignment_id = assignment_id
def on_server_death(self):
self.server_died = True
def setUp(self):
self.AGENT_ALIVE_PACKET = Packet(
MESSAGE_ID_1,
data_model.AGENT_ALIVE,
self.SENDER_ID,
self.WORLD_ID,
self.ASSIGNMENT_ID,
self.DATA,
self.CONVERSATION_ID,
)
self.MESSAGE_SEND_PACKET_1 = Packet(
MESSAGE_ID_2,
data_model.WORLD_MESSAGE,
self.WORLD_ID,
self.SENDER_ID,
self.ASSIGNMENT_ID,
self.DATA,
self.CONVERSATION_ID,
)
self.MESSAGE_SEND_PACKET_2 = Packet(
MESSAGE_ID_3,
data_model.MESSAGE_BATCH,
self.WORLD_ID,
self.SENDER_ID,
self.ASSIGNMENT_ID,
self.DATA,
self.CONVERSATION_ID,
)
self.MESSAGE_SEND_PACKET_3 = Packet(
MESSAGE_ID_4,
data_model.MESSAGE_BATCH,
self.WORLD_ID,
self.SENDER_ID,
self.ASSIGNMENT_ID,
self.DATA,
self.CONVERSATION_ID,
)
self.fake_socket = MockSocket()
time.sleep(0.3)
self.alive_packet = None
self.message_packet = None
self.dead_worker_id = None
self.dead_assignment_id = None
self.server_died = False
self.socket_manager = SocketManager(
'https://127.0.0.1',
self.fake_socket.port,
self.on_alive,
self.on_message,
self.on_worker_death,
TASK_GROUP_ID_1,
1,
self.on_server_death,
)
def tearDown(self):
self.socket_manager.shutdown()
self.fake_socket.close()
def test_init_state(self):
"""
Ensure all of the initial state of the socket_manager is ready.
"""
self.assertEqual(self.socket_manager.server_url, 'https://127.0.0.1')
self.assertEqual(self.socket_manager.port, self.fake_socket.port)
self.assertEqual(self.socket_manager.alive_callback, self.on_alive)
self.assertEqual(self.socket_manager.message_callback, self.on_message)
self.assertEqual(self.socket_manager.socket_dead_callback,
self.on_worker_death)
self.assertEqual(self.socket_manager.task_group_id, TASK_GROUP_ID_1)
self.assertEqual(self.socket_manager.missed_pongs,
1 + (1 / SocketManager.PING_RATE))
self.assertIsNotNone(self.socket_manager.ws)
self.assertTrue(self.socket_manager.keep_running)
self.assertIsNotNone(self.socket_manager.listen_thread)
self.assertSetEqual(self.socket_manager.open_channels, set())
self.assertDictEqual(self.socket_manager.packet_map, {})
self.assertTrue(self.socket_manager.alive)
self.assertFalse(self.socket_manager.is_shutdown)
self.assertEqual(self.socket_manager.get_my_sender_id(), self.WORLD_ID)
def _send_packet_in_background(self, packet, send_time):
"""
creates a thread to handle waiting for a packet send.
"""
def do_send():
self.socket_manager._send_packet(packet, send_time)
self.sent = True
send_thread = threading.Thread(target=do_send, daemon=True)
send_thread.start()
time.sleep(0.02)
def test_packet_send(self):
"""
Checks to see if packets are working.
"""
self.socket_manager._safe_send = mock.MagicMock()
self.sent = False
# Test a blocking acknowledged packet
send_time = time.time()
self.assertEqual(self.MESSAGE_SEND_PACKET_2.status, Packet.STATUS_INIT)
self._send_packet_in_background(self.MESSAGE_SEND_PACKET_2, send_time)
self.assertEqual(self.MESSAGE_SEND_PACKET_2.status, Packet.STATUS_SENT)
self.socket_manager._safe_send.assert_called_once()
self.assertTrue(self.sent)
used_packet_json = self.socket_manager._safe_send.call_args[0][0]
used_packet_dict = json.loads(used_packet_json)
self.assertEqual(used_packet_dict['type'], data_model.MESSAGE_BATCH)
self.assertDictEqual(used_packet_dict['content'],
self.MESSAGE_SEND_PACKET_2.as_dict())
def test_simple_packet_channel_management(self):
"""
Ensure that channels are created, managed, and then removed as expected.
"""
use_packet = self.MESSAGE_SEND_PACKET_1
worker_id = use_packet.receiver_id
assignment_id = use_packet.assignment_id
# Open a channel and assert it is there
self.socket_manager.open_channel(worker_id, assignment_id)
time.sleep(0.1)
connection_id = use_packet.get_receiver_connection_id()
self.assertIn(connection_id, self.socket_manager.open_channels)
self.assertTrue(self.socket_manager.socket_is_open(connection_id))
self.assertFalse(self.socket_manager.socket_is_open(FAKE_ID))
# Send a packet to an open socket, ensure it got queued
resp = self.socket_manager.queue_packet(use_packet)
self.assertIn(use_packet.id, self.socket_manager.packet_map)
self.assertTrue(resp)
# Assert we can get the status of a packet in the map, but not
# existing doesn't throw an error
self.assertEqual(self.socket_manager.get_status(use_packet.id),
use_packet.status)
self.assertEqual(self.socket_manager.get_status(FAKE_ID),
Packet.STATUS_NONE)
# Assert that closing a thread does the correct cleanup work
self.socket_manager.close_channel(connection_id)
time.sleep(0.2)
self.assertNotIn(connection_id, self.socket_manager.open_channels)
self.assertNotIn(use_packet.id, self.socket_manager.packet_map)
# Assert that opening multiple and closing them is possible
self.socket_manager.open_channel(worker_id, assignment_id)
self.socket_manager.open_channel(worker_id + '2', assignment_id)
time.sleep(0.1)
self.assertEqual(len(self.socket_manager.open_channels), 2)
self.socket_manager.close_all_channels()
time.sleep(0.1)
self.assertEqual(len(self.socket_manager.open_channels), 0)