def test_localtaskjob_maintain_heart_rate(self): dagbag = models.DagBag( dag_folder=TEST_DAG_FOLDER, include_examples=False, ) dag = dagbag.dags.get('test_localtaskjob_double_trigger') task = dag.get_task('test_localtaskjob_double_trigger_task') session = settings.Session() dag.clear() dag.create_dagrun(run_id="test", state=State.SUCCESS, execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE, session=session) ti_run = TI(task=task, execution_date=DEFAULT_DATE) ti_run.refresh_from_db() job1 = LocalTaskJob(task_instance=ti_run, executor=SequentialExecutor()) # this should make sure we only heartbeat once and exit at the second # loop in _execute() return_codes = [None, 0] def multi_return_code(): return return_codes.pop(0) time_start = time.time() from airflow.task.task_runner.standard_task_runner import StandardTaskRunner with patch.object(StandardTaskRunner, 'start', return_value=None) as mock_start: with patch.object(StandardTaskRunner, 'return_code') as mock_ret_code: mock_ret_code.side_effect = multi_return_code job1.run() self.assertEqual(mock_start.call_count, 1) self.assertEqual(mock_ret_code.call_count, 2) time_end = time.time() self.assertEqual(self.mock_base_job_sleep.call_count, 1) self.assertEqual(job1.state, State.SUCCESS) # Consider we have patched sleep call, it should not be sleeping to # keep up with the heart rate in other unpatched places # # We already make sure patched sleep call is only called once self.assertLess(time_end - time_start, job1.heartrate) session.close()
def test_loading_local_settings_without_logging_config(self): from airflow.logging_config import configure_logging, log with patch.object(log, 'debug') as mock_info: configure_logging() mock_info.assert_called_with( 'Unable to load custom logging, using default config instead' )
def test_when_the_config_key_does_not_exists(self): from airflow import logging_config with conf_vars({('core', 'logging_config_class'): None}): with patch.object(logging_config.log, 'debug') as mock_debug: logging_config.configure_logging() mock_debug.assert_any_call( 'Could not find key logging_config_class in config')
def test_loading_local_settings_without_logging_config(self): from airflow.logging_config import configure_logging, log with patch.object(log, 'debug') as mock_info: configure_logging() mock_info.assert_called_with( 'Unable to load custom logging, using default config instead' )
def test_loading_valid_local_settings(self): with settings_context(SETTINGS_FILE_VALID): from airflow.logging_config import configure_logging, log with patch.object(log, 'info') as mock_info: configure_logging() mock_info.assert_called_with( 'Successfully imported user-defined logging config from %s', '{}.LOGGING_CONFIG'.format(SETTINGS_DEFAULT_NAME))
def test_loading_invalid_local_settings(self): from airflow.logging_config import configure_logging, log with settings_context(SETTINGS_FILE_INVALID): with patch.object(log, 'warning') as mock_info: # Load config with self.assertRaises(ValueError): configure_logging() mock_info.assert_called_with( 'Unable to load the config, contains a configuration error.' )
def test_loading_valid_local_settings(self): with settings_context(SETTINGS_FILE_VALID): from airflow.logging_config import configure_logging, log with patch.object(log, 'info') as mock_info: configure_logging() mock_info.assert_called_with( 'Successfully imported user-defined logging config from %s', '{}.LOGGING_CONFIG'.format( SETTINGS_DEFAULT_NAME ) )
def test_loading_invalid_local_settings(self): from airflow.logging_config import configure_logging, log with settings_context(SETTINGS_FILE_INVALID): with patch.object(log, 'warning') as mock_info: # Load config with self.assertRaises(ValueError): configure_logging() mock_info.assert_called_with( 'Unable to load the config, contains a configuration error.' )
def test_loading_valid_complex_local_settings(self): # Test what happens when the config is somewhere in a subfolder module_structure = 'etc.airflow.config' dir_structure = module_structure.replace('.', '/') with settings_context(SETTINGS_FILE_VALID, dir_structure): from airflow.logging_config import configure_logging, log with patch.object(log, 'info') as mock_info: configure_logging() mock_info.assert_called_with( 'Successfully imported user-defined logging config from %s', 'etc.airflow.config.{}.LOGGING_CONFIG'.format( SETTINGS_DEFAULT_NAME))
def test_loading_valid_complex_local_settings(self): # Test what happens when the config is somewhere in a subfolder module_structure = 'etc.airflow.config' dir_structure = module_structure.replace('.', '/') with settings_context(SETTINGS_FILE_VALID, dir_structure): from airflow.logging_config import configure_logging, log with patch.object(log, 'info') as mock_info: configure_logging() mock_info.assert_called_with( 'Successfully imported user-defined logging config from %s', 'etc.airflow.config.{}.LOGGING_CONFIG'.format( SETTINGS_DEFAULT_NAME ) )
def test_volume_mount(self): with patch.object(PodLauncher, 'log') as mock_logger: volume_mount = VolumeMount('test-volume', mount_path='/tmp/test_volume', sub_path=None, read_only=False) volume_config = { 'persistentVolumeClaim': { 'claimName': 'test-volume' } } volume = Volume(name='test-volume', configs=volume_config) args = [ "echo \"retrieved from mount\" > /tmp/test_volume/test.txt " "&& cat /tmp/test_volume/test.txt" ] k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=args, labels={"foo": "bar"}, volume_mounts=[volume_mount], volumes=[volume], name="test", task_id="task", in_cluster=False, do_xcom_push=False, ) context = create_context(k) k.execute(context=context) mock_logger.info.assert_any_call(b"retrieved from mount\n") actual_pod = self.api_client.sanitize_for_serialization(k.pod) self.expected_pod['spec']['containers'][0]['args'] = args self.expected_pod['spec']['containers'][0]['volumeMounts'] = [{ 'name': 'test-volume', 'mountPath': '/tmp/test_volume', 'readOnly': False }] self.expected_pod['spec']['volumes'] = [{ 'name': 'test-volume', 'persistentVolumeClaim': { 'claimName': 'test-volume' } }] self.assertEqual(self.expected_pod, actual_pod)
def test_execute_with_empty_destination_bucket(self, mock_hook): mock_hook.return_value.list.return_value = SOURCE_FILES_LIST operator = GoogleCloudStorageToGoogleCloudStorageOperator( task_id=TASK_ID, source_bucket=TEST_BUCKET, source_object=SOURCE_OBJECT_NO_WILDCARD, destination_bucket=None, destination_object=DESTINATION_OBJECT_PREFIX) with patch.object(operator.log, 'warning') as mock_warn: operator.execute(None) mock_warn.assert_called_with( 'destination_bucket is None. Defaulting it to source_bucket (%s)', TEST_BUCKET ) self.assertEquals(operator.destination_bucket, operator.source_bucket)
def test_execute_with_empty_destination_bucket(self, mock_hook): mock_hook.return_value.list.return_value = SOURCE_FILES_LIST operator = GoogleCloudStorageToGoogleCloudStorageOperator( task_id=TASK_ID, source_bucket=TEST_BUCKET, source_object=SOURCE_OBJECT_NO_WILDCARD, destination_bucket=None, destination_object=DESTINATION_OBJECT_PREFIX) with patch.object(operator.log, 'warning') as mock_warn: operator.execute(None) mock_warn.assert_called_with( 'destination_bucket is None. Defaulting it to source_bucket (%s)', TEST_BUCKET ) self.assertEquals(operator.destination_bucket, operator.source_bucket)
def test_when_the_config_key_does_not_exists(self): from airflow import logging_config conf_get = conf.get def side_effect(*args): if args[1] == 'logging_config_class': raise AirflowConfigException else: return conf_get(*args) logging_config.conf.get = mock.Mock(side_effect=side_effect) with patch.object(logging_config.log, 'debug') as mock_debug: logging_config.configure_logging() mock_debug.assert_any_call( 'Could not find key logging_config_class in config')
def test_when_the_config_key_does_not_exists(self): from airflow import logging_config conf_get = conf.get def side_effect(*args): if args[1] == 'logging_config_class': raise AirflowConfigException else: return conf_get(*args) logging_config.conf.get = mock.Mock(side_effect=side_effect) with patch.object(logging_config.log, 'debug') as mock_debug: logging_config.configure_logging() mock_debug.assert_any_call( 'Could not find key logging_config_class in config' )
def test_localtaskjob_double_trigger(self): dagbag = models.DagBag( dag_folder=TEST_DAG_FOLDER, include_examples=False, ) dag = dagbag.dags.get('test_localtaskjob_double_trigger') task = dag.get_task('test_localtaskjob_double_trigger_task') session = settings.Session() dag.clear() dr = dag.create_dagrun(run_id="test", state=State.SUCCESS, execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE, session=session) ti = dr.get_task_instance(task_id=task.task_id, session=session) ti.state = State.RUNNING ti.hostname = get_hostname() ti.pid = 1 session.merge(ti) session.commit() ti_run = TI(task=task, execution_date=DEFAULT_DATE) ti_run.refresh_from_db() job1 = LocalTaskJob(task_instance=ti_run, executor=SequentialExecutor()) from airflow.task.task_runner.standard_task_runner import StandardTaskRunner with patch.object(StandardTaskRunner, 'start', return_value=None) as mock_method: job1.run() mock_method.assert_not_called() ti = dr.get_task_instance(task_id=task.task_id, session=session) self.assertEqual(ti.pid, 1) self.assertEqual(ti.state, State.RUNNING) session.close()