def get_test_agent(db: MephistoDB, unit_id=None) -> str:
# Check creation and retrieval of a agent
worker_name, worker_id = get_test_worker(db)
if unit_id is None:
unit_id = get_test_unit(db)
provider_type = "mock"
task_type = "mock"
unit = Unit(db, unit_id)
return db.new_agent(
worker_id,
unit.db_id,
unit.task_id,
unit.task_run_id,
unit.assignment_id,
unit.task_type,
unit.provider_type,
)
def make_completed_unit(db: MephistoDB) -> str:
"""
Creates a completed unit for the most recently created task run
using some worker. Assumes
"""
workers = db.find_workers()
assert len(workers) > 0, "Must have at least one worker in database"
worker = workers[-1]
task_runs = db.find_task_runs(is_completed=False)
assert len(task_runs) > 0, "Must be at least one incomplete task run"
task_run = task_runs[-1]
assign_id = db.new_assignment(
task_run.task_id,
task_run.db_id,
task_run.requester_id,
task_run.task_type,
task_run.provider_type,
)
unit_id = db.new_unit(
task_run.task_id,
task_run.db_id,
task_run.requester_id,
assign_id,
0,
0.2,
task_run.provider_type,
task_run.task_type,
)
agent_id = db.new_agent(
worker.db_id,
unit_id,
task_run.task_id,
task_run.db_id,
assign_id,
task_run.task_type,
task_run.provider_type,
)
agent = Agent(db, agent_id)
agent.mark_done()
unit = Unit(db, unit_id)
unit.sync_status()
return unit.db_id