async def test_listen_comp_tasks_task( mock_project_subsystem: Dict, comp_task_listening_task: asyncio.Task, client, upd_value: Dict[str, Any], exp_calls: List[str], task_class: NodeClass, ): db_engine: aiopg.sa.Engine = client.app[APP_DB_ENGINE_KEY] async with db_engine.acquire() as conn: # let's put some stuff in there now result = await conn.execute( comp_tasks.insert() .values(outputs=json.dumps({}), node_class=task_class) .returning(literal_column("*")) ) row: RowProxy = await result.fetchone() task = dict(row) # let's update some values await conn.execute( comp_tasks.update() .values(**upd_value) .where(comp_tasks.c.task_id == task["task_id"]) ) for key, mock_fct in mock_project_subsystem.items(): if key in exp_calls: await _wait_for_call(mock_fct) else: mock_fct.assert_not_called()
async def test_listen_comp_tasks_task( mock_project_subsystem: Dict, comp_task_listening_task: asyncio.Task, client, update_values: Dict[str, Any], expected_calls: List[str], task_class: NodeClass, ): db_engine: aiopg.sa.Engine = client.app[APP_DB_ENGINE_KEY] async with db_engine.acquire() as conn: # let's put some stuff in there now result = await conn.execute(comp_tasks.insert().values( outputs=json.dumps({}), node_class=task_class).returning(literal_column("*"))) row: RowProxy = await result.fetchone() task = dict(row) # let's update some values await conn.execute(comp_tasks.update().values(**update_values).where( comp_tasks.c.task_id == task["task_id"])) # tests whether listener gets hooked calls executed for call_name, mocked_call in mock_project_subsystem.items(): if call_name in expected_calls: async for attempt in _async_retry_if_fails(): with attempt: mocked_call.assert_awaited() else: mocked_call.assert_not_called()
def creator(user: Dict[str, Any], project: ProjectAtDB, **overrides_kwargs) -> List[CompTaskAtDB]: created_tasks: List[CompTaskAtDB] = [] for internal_id, (node_id, node_data) in enumerate(project.workbench.items()): task_config = { "project_id": f"{project.uuid}", "node_id": f"{node_id}", "schema": { "inputs": {}, "outputs": {} }, "inputs": { key: json.loads(value.json(by_alias=True, exclude_unset=True)) if isinstance(value, BaseModel) else value for key, value in node_data.inputs.items() } if node_data.inputs else {}, "outputs": { key: json.loads(value.json(by_alias=True, exclude_unset=True)) if isinstance(value, BaseModel) else value for key, value in node_data.outputs.items() } if node_data.outputs else {}, "image": Image(name=node_data.key, tag=node_data.version).dict( # type: ignore by_alias=True, exclude_unset=True), # type: ignore "node_class": to_node_class(node_data.key), "internal_id": internal_id + 1, "submit": datetime.utcnow(), "job_id": generate_dask_job_id( service_key=node_data.key, service_version=node_data.version, user_id=user["id"], project_id=project.uuid, node_id=NodeID(node_id), ), } task_config.update(**overrides_kwargs) with postgres_db.connect() as conn: result = conn.execute(comp_tasks.insert().values( **task_config).returning(sa.literal_column("*"))) new_task = CompTaskAtDB.parse_obj(result.first()) created_tasks.append(new_task) created_task_ids.extend( [t.task_id for t in created_tasks if t.task_id]) return created_tasks
def creator(project_id: str, node_uuid: str, **overrides) -> str: task_config = { "project_id": project_id, "node_id": node_uuid, } task_config.update(**overrides) with postgres_db.connect() as conn: result = conn.execute(comp_tasks.insert() # pylint: disable=no-value-for-parameter .values(**task_config).returning( comp_tasks.c.task_id)) new_task_id = result.first()[comp_tasks.c.task_id] created_task_ids.append(new_task_id) return node_uuid
async def task( db_connection: SAConnection, db_notification_queue: asyncio.Queue, task_class: NodeClass, ) -> Dict: result = await db_connection.execute(comp_tasks.insert().values( outputs=json.dumps({}), node_class=task_class).returning(literal_column("*"))) row: RowProxy = await result.fetchone() task = dict(row) assert (db_notification_queue.empty( )), "database triggered change although it should only trigger on updates!" yield task
async def _create( tasks: Dict[str, Any], dag: Dict[str, List[str]], inputs: Dict[str, Dict[str, Any]], ) -> str: # add a pipeline with postgres_db.connect() as conn: conn.execute( comp_pipeline.insert().values( # pylint: disable=no-value-for-parameter project_id=project_id, dag_adjacency_list=dag ) ) # create the tasks for each pipeline's node for node_uuid, service in tasks.items(): node_inputs = inputs[node_uuid] conn.execute( comp_tasks.insert().values( # pylint: disable=no-value-for-parameter project_id=project_id, node_id=node_uuid, schema=service["schema"], image=service["image"], inputs=node_inputs, state="PENDING", outputs={}, ) ) # check if file must be uploaded # create the tasks for each pipeline's node for node_uuid, service in tasks.items(): node_inputs = inputs[node_uuid] for input_key in node_inputs: if ( isinstance(node_inputs[input_key], dict) and "path" in node_inputs[input_key] ): # update the files in mock_dir to S3 print("--" * 10) print_module_variables(module=np.node_config) print("--" * 10) PORTS = await np.ports(user_id, project_id, node_uuid) await (await PORTS.inputs)[input_key].set( mock_dir / node_inputs[input_key]["path"] ) return project_id
async def test_listen_comp_tasks_task( mock_project_subsystem: Dict, comp_task_listening_task: None, client, update_values: Dict[str, Any], expected_calls: List[str], task_class: NodeClass, ): db_engine: aiopg.sa.Engine = client.app[APP_DB_ENGINE_KEY] async with db_engine.acquire() as conn: # let's put some stuff in there now result = await conn.execute( comp_tasks.insert() .values(outputs=json.dumps({}), node_class=task_class) .returning(literal_column("*")) ) row: RowProxy = await result.fetchone() task = dict(row) # let's update some values await conn.execute( comp_tasks.update() .values(**update_values) .where(comp_tasks.c.task_id == task["task_id"]) ) # tests whether listener gets hooked calls executed for call_name, mocked_call in mock_project_subsystem.items(): if call_name in expected_calls: async for attempt in AsyncRetrying( wait=wait_fixed(1), stop=stop_after_delay(10), retry=retry_if_exception_type(AssertionError), before_sleep=before_sleep_log(logger, logging.INFO), reraise=True, ): with attempt: mocked_call.assert_awaited() else: mocked_call.assert_not_called()