def get_test_task_run(db: MephistoDB) -> str: """Helper to create a task run for tests""" task_name, task_id = get_test_task(db) requester_name, requester_id = get_test_requester(db) init_params = OmegaConf.to_yaml(OmegaConf.structured(MOCK_CONFIG)) return db.new_task_run(task_id, requester_id, json.dumps(init_params), "mock", "mock")
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
def get_test_assignment(db: MephistoDB) -> str: """Helper to create an assignment for tests""" task_run_id = get_test_task_run(db) task_run = TaskRun(db, task_run_id) return db.new_assignment( task_run.task_id, task_run_id, task_run.requester_id, task_run.task_type, task_run.provider_type, )
def get_test_unit(db: MephistoDB, unit_index=0) -> str: # Check creation and retrieval of a unit assignment_id = get_test_assignment(db) pay_amount = 15.0 assignment = Assignment(db, assignment_id) return db.new_unit( assignment.task_id, assignment.task_run_id, assignment.requester_id, assignment.db_id, 0, pay_amount, assignment.provider_type, assignment.task_type, )
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 get_test_task_run(db: MephistoDB) -> str: """Helper to create a task run for tests""" task_name, task_id = get_test_task(db) requester_name, requester_id = get_test_requester(db) init_params = TaskConfig.get_mock_params() return db.new_task_run(task_id, requester_id, init_params, "mock", "mock")
def get_test_task(db: MephistoDB) -> Tuple[str, str]: """Helper to create a task for tests""" task_name = "test_task" task_type = "mock" task_id = db.new_task(task_name, task_type) return task_name, task_id
def get_test_requester(db: MephistoDB) -> Tuple[str, str]: """Helper to create a requester for tests""" requester_name = "test_requester" provider_type = "mock" requester_id = db.new_requester(requester_name, provider_type) return requester_name, requester_id
def get_test_worker(db: MephistoDB) -> Tuple[str, str]: """Helper to create a worker for tests""" worker_name = "test_worker" provider_type = "mock" worker_id = db.new_worker(worker_name, provider_type) return worker_name, worker_id
def get_test_project(db: MephistoDB) -> Tuple[str, str]: """Helper to create a project for tests""" project_name = "test_project" project_id = db.new_project(project_name) return project_name, project_id