def get_db_status(self) -> str:
"""
Return the status as currently stored in the database
"""
if self.db_status in AssignmentState.final_unit():
return self.db_status
row = self.db.get_unit(self.db_id)
assert row is not None, f"Unit {self.db_id} stopped existing in the db..."
return row["status"]
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
# TODO(#97) add test to ensure this behavior/assumption holds always
if self.db_status in AssignmentState.final_unit():
if self.agent_id is None:
return None
return Agent(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
# TODO(#101) this may not be particularly efficient
row = self.db.get_unit(self.db_id)
assert row is not None, f"Unit {self.db_id} stopped existing in the db..."
agent_id = row["agent_id"]
if agent_id is not None:
return Agent(self.db, agent_id)
return None