def test_job_exe_clean_task(self): """Tests the NodeManager where a cleanup task is returned to clean up a job execution""" when = now() node_mgr = NodeManager() node_mgr.register_agents([self.agent_1, self.agent_2]) node_mgr.sync_with_database(scheduler_mgr.config) cleanup_mgr = CleanupManager() cleanup_mgr.update_nodes(node_mgr.get_nodes()) tasks = node_mgr.get_next_tasks(when) task_mgr = TaskManager() # Complete initial cleanup tasks for task in tasks: task_mgr.launch_tasks([task], now()) update = job_test_utils.create_task_status_update( task.id, task.agent_id, TaskStatusUpdate.FINISHED, now()) task_mgr.handle_task_update(update) node_mgr.handle_task_update(update) # Mark image pull done to get rid of image tasks for node in node_mgr.get_nodes(): node._image_pull_completed() node._update_state() job_exe = job_test_utils.create_running_job_exe(agent_id=self.agent_1, node=self.node_1) # Add a job execution to clean up and get the cleanup task for it cleanup_mgr.add_job_execution(job_exe) tasks = node_mgr.get_next_tasks(when) self.assertEqual(len(tasks), 1) task = tasks[0] self.assertEqual(task.agent_id, self.agent_1.agent_id) self.assertFalse(task.is_initial_cleanup) self.assertEqual(len(task.job_exes), 1)
def test_pull_task_change_agent_id(self): """Tests the NodeManager where a node's agent ID changes during a pull task""" when = now() manager = NodeManager() manager.register_agents([self.agent_1, self.agent_2]) manager.sync_with_database(scheduler_mgr.config) for node in manager.get_nodes(): node._last_health_task = when node._initial_cleanup_completed() node._update_state() tasks = manager.get_next_tasks(when) task_mgr = TaskManager() task_2 = None for task in tasks: task_mgr.launch_tasks([task], when) if task.agent_id == self.agent_2.agent_id: task_2 = task # Node 2 changes agent ID to 3 manager.lost_node(self.agent_2.agent_id) manager.register_agents([self.agent_3]) manager.sync_with_database(scheduler_mgr.config) for node in manager.get_nodes(): node._last_health_task = when node._initial_cleanup_completed() node._update_state() # Should get new Docker pull task for node 2 tasks = manager.get_next_tasks(when) self.assertEqual(len(tasks), 1) new_task_2 = tasks[0] self.assertEqual(new_task_2.agent_id, self.agent_3.agent_id) # Task update comes back for original node 2 Docker pull task, manager should ignore with no exception update = job_test_utils.create_task_status_update( task_2.id, task_2.agent_id, TaskStatusUpdate.FAILED, when) task_mgr.handle_task_update(update) manager.handle_task_update(update)