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_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)
def test_onboarding_agents(self) -> None:
"""Ensure that the db can create and manipulate onboarding agents"""
assert self.db is not None, "No db initialized"
db: MephistoDB = self.db
task_run_id = get_test_task_run(db)
task_run = TaskRun(db, task_run_id)
task = task_run.get_task()
worker_name, worker_id = get_test_worker(db)
onboarding_agent_id = db.new_onboarding_agent(worker_id, task.db_id,
task_run_id, "mock")
self.assertIsNotNone(onboarding_agent_id)
onboarding_agent = OnboardingAgent(db, onboarding_agent_id)
self.assertIsInstance(onboarding_agent, OnboardingAgent)
found_agents = db.find_onboarding_agents(worker_id=worker_id)
self.assertEqual(len(found_agents), 1)
self.assertIsInstance(found_agents[0], OnboardingAgent)
found_agent = found_agents[0]
self.assertEqual(found_agent.db_id, onboarding_agent_id)
self.assertEqual(found_agent.get_status(), AgentState.STATUS_NONE)
def test_qualifications(self) -> None:
"""Test creating, assigning, revoking, and deleting qualifications"""
assert self.db is not None, "No db initialized"
db: MephistoDB = self.db
qualification_name = "TEST_QUALIFICATION_1"
# Create qualification
qual_id = db.make_qualification(qualification_name)
# ensure qualification has been made
qualifications = db.find_qualifications(qualification_name=qualification_name)
self.assertEqual(len(qualifications), 1, "Single qualification not created")
self.assertIsInstance(qualifications[0], Qualification)
# Can't create same qualification again
with self.assertRaises(EntryAlreadyExistsException):
qual_id = db.make_qualification(qualification_name)
qualifications = db.find_qualifications(qualification_name)
self.assertEqual(len(qualifications), 1, "More than one qualification created")
# Grant the qualification to a worker
worker_name, worker_id = get_test_worker(db)
db.grant_qualification(qual_id, worker_id)
# Ensure it was granted
granted_quals = db.check_granted_qualifications()
self.assertEqual(len(granted_quals), 1, "Single qualification not granted")
granted_qual = granted_quals[0]
self.assertEqual(granted_qual.worker_id, worker_id)
self.assertEqual(granted_qual.qualification_id, qual_id)
self.assertEqual(granted_qual.value, 1)
# Update the qualification
db.grant_qualification(qual_id, worker_id, value=3)
# Ensure it was updated
granted_quals = db.check_granted_qualifications()
self.assertEqual(len(granted_quals), 1, "Single qualification not granted")
granted_qual = granted_quals[0]
self.assertEqual(granted_qual.worker_id, worker_id)
self.assertEqual(granted_qual.qualification_id, qual_id)
self.assertEqual(granted_qual.value, 3)
# Delete the qualification
db.revoke_qualification(qual_id, worker_id)
granted_quals = db.check_granted_qualifications()
self.assertEqual(len(granted_quals), 0, "Single qualification not removed")
# Re-grant the qualification
db.grant_qualification(qual_id, worker_id)
# Delete the qualification entirely
db.delete_qualification(qualification_name)
# Ensure deleted and cascaded
qualifications = db.find_qualifications(qualification_name)
self.assertEqual(len(qualifications), 0, "Qualification not remove")
granted_quals = db.check_granted_qualifications()
self.assertEqual(
len(granted_quals), 0, "Cascade granted qualification not removed"
)
# cant retrieve the qualification directly anymore
with self.assertRaises(EntryDoesNotExistException):
qualification_row = db.get_granted_qualification(qual_id, worker_id)