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()