def test_worker_started_with_required_arguments(self, mock_worker, mock_popen, mock_locations):
pid_file = "pid_file"
mock_locations.return_value = (pid_file, None, None, None)
concurrency = '1'
celery_hostname = "celery_hostname"
queues = "queue"
autoscale = "2,5"
args = self.parser.parse_args([
'celery',
'worker',
'--autoscale',
autoscale,
'--concurrency',
concurrency,
'--celery-hostname',
celery_hostname,
'--queues',
queues
])
with mock.patch('celery.platforms.check_privileges') as mock_privil:
mock_privil.return_value = 0
celery_command.worker(args)
mock_worker.worker.return_value.run.assert_called_once_with(
pool='prefork',
optimization='fair',
O='fair', # noqa
queues=queues,
pidfile=pid_file,
concurrency=int(concurrency),
autoscale=autoscale,
hostname=celery_hostname,
loglevel=conf.get('logging', 'LOGGING_LEVEL'),
)
def test_skip_serve_logs_on_worker_start(self, mock_worker):
with mock.patch('airflow.cli.commands.celery_command.Process') as mock_popen:
args = self.parser.parse_args(['celery', 'worker', '--concurrency', '1', '--skip-serve-logs'])
with mock.patch('celery.platforms.check_privileges') as mock_privil:
mock_privil.return_value = 0
celery_command.worker(args)
mock_popen.assert_not_called()
def test_error(self, mock_validate_session):
"""
Test to verify the exit mechanism of airflow-worker cli
by mocking validate_session method
"""
mock_validate_session.return_value = False
with self.assertRaises(SystemExit) as cm:
celery_command.worker(mock_args)
self.assertEqual(cm.exception.code, 1)
def test_serve_logs_on_worker_start(self):
with mock.patch(
'airflow.cli.commands.celery_command.Process') as mock_process:
args = self.parser.parse_args(['celery', 'worker', '-c', '-1'])
with mock.patch(
'celery.platforms.check_privileges') as mock_privil:
mock_privil.return_value = 0
celery_command.worker(args)
mock_process.assert_called()
def test_error(self, mock_validate_session):
"""
Test to verify the exit mechanism of airflow-worker cli
by mocking validate_session method
"""
mock_validate_session.return_value = False
with pytest.raises(SystemExit) as ctx:
celery_command.worker(Namespace(queues=1, concurrency=1))
assert str(ctx.value
) == "Worker exiting, database connection precheck failed."
def test_same_pid_file_is_used_in_start_and_stop(
self, mock_setup_locations, mock_celery_worker,
mock_read_pid_from_pidfile):
pid_file = "test_pid_file"
mock_setup_locations.return_value = (pid_file, None, None, None)
mock_read_pid_from_pidfile.return_value = None
# Call worker
worker_args = self.parser.parse_args(
['celery', 'worker', '--skip-serve-logs'])
celery_command.worker(worker_args)
assert mock_celery_worker.call_args
assert 'pidfile' in mock_celery_worker.call_args.kwargs
assert mock_celery_worker.call_args.kwargs['pidfile'] == pid_file
# Call stop
stop_args = self.parser.parse_args(['celery', 'stop'])
celery_command.stop_worker(stop_args)
mock_read_pid_from_pidfile.assert_called_once_with(pid_file)