def _new_agent( self, worker_id: str, unit_id: str, task_id: str, task_run_id: str, assignment_id: str, task_type: str, provider_type: str, ) -> str: """ Wrapper around the new_agent call that finds and updates the unit locally too, as this isn't guaranteed otherwise but is an important part of the singleton """ agent_id = super()._new_agent( worker_id, unit_id, task_id, task_run_id, assignment_id, task_type, provider_type, ) agent = Agent.get(self, agent_id) unit = agent.get_unit() unit.agent_id = agent_id unit.db_status = AssignmentState.ASSIGNED unit.worker_id = agent.worker_id return agent_id
def get_assigned_agent(self) -> Optional[Agent]: """ Get the agent assigned to this Unit if there is one, else return None """ # In these statuses, we know the agent isn't changing anymore, and thus will # not need to be re-queried if self.db_status in AssignmentState.final_unit(): if self.agent_id is None: return None return Agent.get(self.db, self.agent_id) # Query the database to get the most up-to-date assignment, as this can # change after instantiation if the Unit status isn't final unit_copy = Unit.get(self.db, self.db_id) self.agent_id = unit_copy.agent_id if self.agent_id is not None: return Agent.get(self.db, self.agent_id) return None
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.get(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.get(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_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.get(db, self.get_fake_id("Agent")) unit_id = get_test_unit(db) worker_name, worker_id = get_test_worker(db) unit = Unit.get(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)