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)