def test_agent(self) -> None: """Test creation and querying of agents""" assert self.db is not None, "No db initialized" db: MephistoDB = self.db # Check creation and retrieval of a agent worker_name, worker_id = get_test_worker(db) unit_id = get_test_unit(db) unit = Unit(db, unit_id) agent_id = db.new_agent( worker_id, unit_id, unit.task_id, unit.task_run_id, unit.assignment_id, unit.task_type, unit.provider_type, ) self.assertIsNotNone(agent_id) self.assertTrue(isinstance(agent_id, str)) agent_row = db.get_agent(agent_id) self.assertEqual(agent_row["worker_id"], worker_id) self.assertEqual(agent_row["unit_id"], unit_id) self.assertEqual(agent_row["status"], AgentState.STATUS_NONE) # ensure the unit is assigned now units = db.find_units(status=AssignmentState.ASSIGNED) self.assertEqual(len(units), 1) agent = Agent(db, agent_id) self.assertEqual(agent.worker_id, worker_id) # Check finding for agents agents = db.find_agents() self.assertEqual(len(agents), 1) self.assertTrue(isinstance(agents[0], Agent)) self.assertEqual(agents[0].db_id, agent_id) self.assertEqual(agents[0].worker_id, worker_id) # Check finding for specific agents agents = db.find_agents(worker_id=worker_id) self.assertEqual(len(agents), 1) self.assertTrue(isinstance(agents[0], Agent)) self.assertEqual(agents[0].db_id, agent_id) self.assertEqual(agents[0].worker_id, worker_id) agents = db.find_agents(worker_id=self.get_fake_id("Worker")) self.assertEqual(len(agents), 0)
def test_unit_updates(self) -> None: """Test updating a unit's status""" assert self.db is not None, "No db initialized" db: MephistoDB = self.db unit_id = get_test_unit(db) # Check finding for specific units units = db.find_units(status=AssignmentState.COMPLETED) self.assertEqual(len(units), 0) db.update_unit(unit_id, status=AssignmentState.COMPLETED) # Check finding for specific units units = db.find_units(status=AssignmentState.COMPLETED) self.assertEqual(len(units), 1) # Can't update with a status that doesn't exist with self.assertRaises(MephistoDBException): db.update_unit(unit_id, status="FAKE_STATUS")
def test_agent_fails(self) -> None: """Ensure agents fail to be created or loaded under failure conditions""" assert self.db is not None, "No db initialized" db: MephistoDB = self.db # Cant get non-existent entry with self.assertRaises(EntryDoesNotExistException): agent = Agent(db, self.get_fake_id("Agent")) unit_id = get_test_unit(db) worker_name, worker_id = get_test_worker(db) unit = Unit(db, unit_id) # Can't use invalid worker id with self.assertRaises(EntryDoesNotExistException): agent_id = db.new_agent( self.get_fake_id("Worker"), unit_id, unit.task_id, unit.task_run_id, unit.assignment_id, unit.task_type, unit.provider_type, ) # Can't use invalid unit id with self.assertRaises(EntryDoesNotExistException): agent_id = db.new_agent( worker_id, self.get_fake_id("Unit"), unit.task_id, unit.task_run_id, unit.assignment_id, unit.task_type, unit.provider_type, ) # Ensure no agents were created agents = db.find_agents() self.assertEqual(len(agents), 0)