async def test_info_task_filtered_response(can_debug_amqp, container_requester,
dummy_request):
async with container_requester as requester:
task_vars.request.set(dummy_request)
task_vars.db.set(requester.db)
await get_container(requester=requester)
# Add some tasks first
t1 = await add_task(_test_func, 1, 2)
# Update task with job data
state_manager = get_state_manager()
await state_manager.update(t1.task_id, {"job_data": {"foo": "bar"}})
# When the user has debug permission, the response contains job_data
can_debug_amqp.return_value = True
resp, status = await requester(
"GET", f"/db/guillotina/@amqp-tasks/{t1.task_id}")
assert status == 200
assert "job_data" in resp
# When the user has not debug permission, the response does not contain job_data
can_debug_amqp.return_value = False
resp, status = await requester(
"GET", f"/db/guillotina/@amqp-tasks/{t1.task_id}")
assert status == 200
assert "job_data" not in resp
async def test_is_canceled_should_return_true_only_on_canceled_tasks(
configured_state_manager, loop):
state_manager = get_state_manager(loop)
await state_manager.cancel('foo')
assert await state_manager.is_canceled('foo')
assert not await state_manager.is_canceled('bar')
await clear_cache(state_manager)
async def test_list_should_yield_all_items(configured_state_manager, loop):
state_manager = get_state_manager(loop)
tasks = set({'t1', 't2', 't3', 't4', 't5'})
for task_id in tasks:
await state_manager.update(task_id, {'state': f'{task_id} is OK!'})
assert set([tid async for tid in state_manager.list()]) == tasks
await clear_cache(state_manager)
async def test_list_should_yield_all_items(configured_state_manager, loop):
state_manager = get_state_manager(loop)
tasks = set({"t1", "t2", "t3", "t4", "t5"})
for task_id in tasks:
await state_manager.update(task_id, {"state": f"{task_id} is OK!"})
assert set([tid async for tid in state_manager.list()]) == tasks
await clear_cache(state_manager)
async def test_task_state_disappears_after_ttl(redis_state_manager, loop):
state_manager = get_state_manager(loop)
await state_manager.update('foo', {'state': 'bar'}, ttl=2)
data = await state_manager.get('foo')
assert data['state'] == 'bar'
await asyncio.sleep(2)
data = await state_manager.get('foo')
assert not data
await clear_cache(state_manager)
async def test_connection_reset_errors_are_retried(redis_state_manager, loop):
state_manager = get_state_manager(loop)
mocked = MockedRedisGET()
with asynctest.mock.patch("guillotina_amqp.state.aioredis.Redis.get", new=mocked):
with pytest.raises(ConnectionResetError):
await state_manager.get("foo")
assert mocked.called == 4
async def test_acquire_should_unblock_after_ttl(configured_state_manager, loop):
state_manager = get_state_manager(loop)
await state_manager.acquire("foo", ttl=1)
with pytest.raises(TaskAlreadyAcquired):
await state_manager.acquire("foo", ttl=300)
# Wait for the ttl and check
await asyncio.sleep(1.1)
await state_manager.acquire("foo", ttl=1)
await clear_cache(state_manager)
async def test_task_state_disappears_after_ttl(redis_state_manager, loop):
state_manager = get_state_manager(loop)
await state_manager.update("foo", {"state": "bar"}, ttl=2)
data = await state_manager.get("foo")
assert data["state"] == "bar"
await asyncio.sleep(2)
data = await state_manager.get("foo")
assert not data
await clear_cache(state_manager)
async def test_acquire_should_block_further_acquires(configured_state_manager, loop):
state_manager = get_state_manager(loop)
FOREVER = 300
await state_manager.acquire("foo", ttl=FOREVER)
with pytest.raises(TaskAlreadyAcquired):
await state_manager.acquire("foo", ttl=FOREVER)
# Release it and acquire again
await state_manager.release("foo")
await state_manager.acquire("foo", ttl=FOREVER)
await clear_cache(state_manager)
async def test_cancel_should_put_tasks_in_cancelation_list(
configured_state_manager, loop):
state_manager = get_state_manager(loop)
await state_manager.update('foo', {'status': 'bar'})
canceled_list = [tid async for tid in state_manager.cancelation_list()]
assert 'foo' not in canceled_list
await state_manager.cancel('foo')
canceled_list = [tid async for tid in state_manager.cancelation_list()]
assert 'foo' in canceled_list
await clear_cache(state_manager)
async def test_clean_cancel_should_clean_from_canceled_list(
configured_state_manager, loop):
state_manager = get_state_manager(loop)
await state_manager.cancel('foo')
canceled_list = [tid async for tid in state_manager.cancelation_list()]
assert 'foo' in canceled_list
await state_manager.clean_canceled('foo')
canceled_list = [tid async for tid in state_manager.cancelation_list()]
assert 'foo' not in canceled_list
await clear_cache(state_manager)
async def test_cancel_should_put_tasks_in_cancelation_list(
configured_state_manager, loop
):
state_manager = get_state_manager(loop)
await state_manager.update("foo", {"status": "bar"})
canceled_list = [tid async for tid in state_manager.cancelation_list()]
assert "foo" not in canceled_list
await state_manager.cancel("foo")
canceled_list = [tid async for tid in state_manager.cancelation_list()]
assert "foo" in canceled_list
await clear_cache(state_manager)
async def test_release_should_raise_if_task_is_not_yours(
configured_state_manager, loop):
state_manager = get_state_manager(loop)
state_manager.worker_id = 'me'
await state_manager.acquire('footask', ttl=120)
with pytest.raises(TaskAccessUnauthorized):
state_manager.worker_id = 'another_person'
await state_manager.release('footask')
state_manager.worker_id = 'me'
await state_manager.release('footask')
await clear_cache(state_manager)
async def test_refresh_should_raise_if_task_is_not_yours(
configured_state_manager, loop
):
state_manager = get_state_manager(loop)
state_manager.worker_id = "me"
await state_manager.acquire("footask", ttl=120)
await state_manager.refresh_lock("footask", ttl=20)
with pytest.raises(TaskAccessUnauthorized):
state_manager.worker_id = "another_person"
await state_manager.refresh_lock("footask", 120)
state_manager.worker_id = "me"
await state_manager.refresh_lock("footask", 20)
await clear_cache(state_manager)
async def test_worker_acks_already_acquired_tasks(dummy_request):
# Fake some task data
task_id = "foo"
task_data = json.dumps({"task_id": task_id, "func": "foo.bar"})
# Mock as if the task would be acquired
state_manager = get_state_manager()
await state_manager.acquire(task_id, 900)
channel = MockChannel()
assert len(channel.acked) == 0
envelope = MockEnvelope("footag")
# Pretend worker picks up the task
worker = Worker()
await worker.handle_queued_job(channel, task_data, envelope, None)
# Check that worker sent ack to amqp channel
assert len(channel.acked) == 1
assert channel.acked[0]["kwargs"]["delivery_tag"] == envelope.delivery_tag
async def wait_for_all_tasks() -> None:
# Wait for all scheduled tasks
sm = get_state_manager()
async for task_id in sm.list():
await wait_for_task(task_id)
async def test_update_and_get_should_match(configured_state_manager, loop):
state_manager = get_state_manager(loop)
await state_manager.update("foo", {"state": "bar"})
data = await state_manager.get("foo")
assert data["state"] == "bar"
await clear_cache(state_manager)
async def add_task(func,
*args,
_request=None,
_retries=3,
_task_id=None,
**kwargs):
"""Given a function and its arguments, it adds it as a task to be ran
by workers.
"""
# Get the request and prepare request data
if _request is None:
_request = get_current_request()
req_data: SerializedRequest = serialize_request(_request)
if _task_id is None:
task_id = generate_task_id()
else:
task_id = _task_id
dotted_name = get_dotted_name(func)
retries = 0
while True:
# Get the rabbitmq connection
try:
channel, transport, protocol = await amqp.get_connection()
except AMQPConfigurationNotFoundError:
logger.warning(
f"Could not schedule {dotted_name}, AMQP settings not configured"
)
return
try:
state = TaskState(task_id)
db = task_vars.db.get()
container = task_vars.container.get()
logger.info(f"Scheduling task: {task_id}: {dotted_name}")
data = json.dumps({
"func": dotted_name,
"args": args,
"kwargs": kwargs,
"db_id": getattr(db, "id", None),
"container_id": getattr(container, "id", None),
"req_data": req_data,
"task_id": task_id,
})
# Publish task data on rabbitmq
await channel.publish(
data,
exchange_name=app_settings["amqp"]["exchange"],
routing_key=app_settings["amqp"]["queue"],
properties={"delivery_mode": 2},
)
# Update tasks's global state
state_manager = get_state_manager()
await update_task_scheduled(state_manager,
task_id,
updated=time.time())
logger.info(f"Scheduled task: {task_id}: {dotted_name}")
return state
except (aioamqp.AmqpClosedConnection,
aioamqp.exceptions.ChannelClosed):
await amqp.remove_connection()
if retries >= _retries:
raise
retries += 1
async def test_update_and_get_should_match(configured_state_manager, loop):
state_manager = get_state_manager(loop)
await state_manager.update('foo', {'state': 'bar'})
data = await state_manager.get('foo')
assert data['state'] == 'bar'
await clear_cache(state_manager)
def state_manager(self):
if self._state_manager is None:
self._state_manager = get_state_manager()
return self._state_manager
async def add_task(func,
*args,
_request=None,
_retries=3,
_task_id=None,
**kwargs):
"""Given a function and its arguments, it adds it as a task to be ran
by workers.
"""
# Get the request and prepare request data
if _request is None:
_request = get_current_request()
req_data = {
'url': str(_request.url),
'headers': dict(_request.headers),
'method': _request.method,
'annotations': getattr(_request, 'annotations', {})
}
user = get_authenticated_user()
if user is not None:
try:
req_data['user'] = {
'id':
user.id,
'roles': [
name for name, setting in user.roles.items()
if setting == Allow
],
'groups':
user.groups,
'headers':
dict(_request.headers),
'data':
getattr(user, 'data', {})
}
except AttributeError:
pass
container = task_vars.container.get()
if container is not None:
req_data['container_url'] = IAbsoluteURL(container, _request)()
if _task_id is None:
task_id = generate_task_id()
else:
task_id = _task_id
retries = 0
while True:
# Get the rabbitmq connection
channel, transport, protocol = await amqp.get_connection()
try:
state = TaskState(task_id)
dotted_name = get_dotted_name(func)
db = task_vars.db.get()
logger.info(f'Scheduling task: {task_id}: {dotted_name}')
data = json.dumps({
'func': dotted_name,
'args': args,
'kwargs': kwargs,
'db_id': getattr(db, 'id', None),
'container_id': getattr(container, 'id', None),
'req_data': req_data,
'task_id': task_id
})
# Publish task data on rabbitmq
await channel.publish(
data,
exchange_name=app_settings['amqp']['exchange'],
routing_key=app_settings['amqp']['queue'],
properties={'delivery_mode': 2})
# Update tasks's global state
state_manager = get_state_manager()
await update_task_scheduled(state_manager,
task_id,
updated=time.time())
logger.info(f'Scheduled task: {task_id}: {dotted_name}')
return state
except (aioamqp.AmqpClosedConnection,
aioamqp.exceptions.ChannelClosed):
await amqp.remove_connection()
if retries >= _retries:
raise
retries += 1
