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