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 write_ports_configuration(
self, json_configuration: str, project_id: str, node_uuid: str
):
message = (
f"Writing port configuration to database for "
f"project={project_id} node={node_uuid}: {json_configuration}"
)
log.debug(message)
node_configuration = json.loads(json_configuration)
async with DBContextManager(self._db_engine) as engine:
async with engine.acquire() as connection:
# update the necessary parts
await connection.execute(
# FIXME: E1120:No value for argument 'dml' in method call
# pylint: disable=E1120
comp_tasks.update()
.where(
and_(
comp_tasks.c.node_id == node_uuid,
comp_tasks.c.project_id == project_id,
)
)
.values(
schema=node_configuration["schema"],
inputs=node_configuration["inputs"],
outputs=node_configuration["outputs"],
run_hash=node_configuration.get("run_hash"),
)
)
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()
async def _set_task_state(
aiopg_engine: Iterator[aiopg.sa.engine.Engine],
node_id: str,
state: StateType # type: ignore
):
async with aiopg_engine.acquire() as conn: # type: ignore
await conn.execute(comp_tasks.update().where(
comp_tasks.c.node_id == node_id).values(state=state))
async def set_comp_task_outputs(
aiopg_engine: aiopg.sa.engine.Engine,
node_id: NodeID,
outputs_schema: Dict[str, Any],
outputs: Dict[str, Any] # type: ignore
):
async with aiopg_engine.acquire() as conn: # type: ignore
await conn.execute(
# pylint: disable=no-value-for-parameter
comp_tasks.update().where(comp_tasks.c.node_id == f"{node_id}"
).values(outputs=outputs,
schema={
"outputs": outputs_schema,
"inputs": {}
}))
async def write_ports_configuration(self, json_configuration: str,
node_uuid: str):
log.debug("Writing ports configuration to database")
node_configuration = json.loads(json_configuration)
async with DBContextManager(self._db_engine) as engine:
async with engine.acquire() as connection:
# update the necessary parts
await connection.execute(
# FIXME: E1120:No value for argument 'dml' in method call
# pylint: disable=E1120
comp_tasks.update().where(
and_(
comp_tasks.c.node_id == node_uuid,
comp_tasks.c.project_id == config.PROJECT_ID,
)).values(
schema=node_configuration["schema"],
inputs=node_configuration["inputs"],
outputs=node_configuration["outputs"],
))
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()
def update_configuration(
postgres_engine: sa.engine.Engine,
project_id: str,
node_uuid: str,
new_configuration: Dict,
) -> None:
log.debug(
"Update configuration of pipeline %s, node %s",
project_id,
node_uuid,
)
with postgres_engine.connect() as conn:
conn.execute(comp_tasks.update() # pylint: disable=no-value-for-parameter
.where((comp_tasks.c.project_id == project_id)
& (comp_tasks.c.node_id == node_uuid)).values(
schema=new_configuration["schema"],
inputs=new_configuration["inputs"],
outputs=new_configuration["outputs"],
))
log.debug("Updated configuration")
async def _update_comp_task_with(conn: SAConnection, task: Dict, **kwargs):
await conn.execute(comp_tasks.update().values(**kwargs).where(
comp_tasks.c.task_id == task["task_id"]))
