def test_scale_celery_task(self, mock_scale_by_runs, mock_scale_by_tasks):
with self.settings(CELERY_SCALE_BY_RUN=True):
scale_celery_task(16000)
mock_scale_by_runs.assert_called_once_with(16000)
with self.settings(CELERY_SCALE_BY_RUN=False):
scale_celery_task(16000)
celery_tasks = get_celery_tasks_scale_by_task()
mock_scale_by_tasks.assert_called_once_with(celery_tasks, 16000)
def test_scale_by_tasks(
self,
mock_get_scale_client,
mock_get_celery_task_details,
mock_order_celery_tasks,
mock_get_all_rabbitmq_objects,
):
mock_scale_client = Mock()
mock_get_scale_client.return_value = mock_scale_client, "app_name"
# Run until queues are empty.
example_memory_used = 2048
example_disk_used = 3072
mock_scale_client.get_running_tasks_memory.return_value = example_memory_used
mock_scale_client.get_running_tasks.return_value = {"pagination": {"total_results": 1}}
example_queues = [{"name": "celery", "messages": 2}]
empty_queues = [{"name": "celery", "messages": 0}]
mock_get_all_rabbitmq_objects.side_effect = [example_queues, empty_queues]
celery_task_details = {"task_counts": {"celery": 1}, "memory": example_memory_used, "disk": example_disk_used}
mock_get_celery_task_details.return_value = celery_task_details
ordered_celery_tasks = OrderedDict(
{
"queue1": {
"command": "celery -A eventkit_cloud worker --loglevel=$LOG_LEVEL -n worker@%h -Q queue1 ",
"disk": 2048,
"memory": 2048,
},
"celery": {
"command": "celery -A eventkit_cloud worker --loglevel=$LOG_LEVEL -n celery@%h -Q celery ",
"disk": 2048,
"memory": 3072,
"limit": 2,
},
}
)
mock_get_scale_client.return_value = mock_scale_client, "app_name"
mock_order_celery_tasks.return_value = ordered_celery_tasks
celery_tasks = get_celery_tasks_scale_by_task()
scale_by_tasks(celery_tasks, 8000)
mock_scale_client.run_task.assert_called_once()
mock_scale_client.reset_mock()
mock_get_all_rabbitmq_objects.side_effect = None
# Run until memory is exhausted.
mock_get_all_rabbitmq_objects.return_value = example_queues
over_memory = 9000
mock_scale_client.get_running_tasks_memory.return_value = over_memory
scale_by_tasks(celery_tasks, 8000)
mock_scale_client.run_task.assert_called_once()
mock_scale_client.reset_mock()
# Don't run if not enough memory.
mock_get_all_rabbitmq_objects.return_value = example_queues
mock_get_celery_task_details.return_value = celery_task_details
scale_by_tasks(celery_tasks, 1000)
mock_scale_client.run_task.assert_not_called()
mock_scale_client.reset_mock()
# # Don't run if task limit is reached.
mock_scale_client.get_running_tasks.return_value = {"pagination": {"total_results": 3}}
scale_by_tasks(celery_tasks, 8000)
mock_scale_client.run_task.assert_not_called()
mock_scale_client.reset_mock()
