def test_initialize_job(self, mock_get_conn, mock_get_or_create_glue_job,
mock_get_job_state):
some_data_path = "s3://glue-datasets/examples/medicare/SampleData.csv"
some_script_arguments = {"--s3_input_data_path": some_data_path}
some_script = "s3:/glue-examples/glue-scripts/sample_aws_glue_job.py"
some_s3_bucket = "my-includes"
mock_get_or_create_glue_job.Name = mock.Mock(Name='aws_test_glue_job')
mock_get_conn.return_value.start_job_run()
mock_job_run_state = mock_get_job_state.return_value
glue_job_hook = AwsGlueJobHook(
job_name='aws_test_glue_job',
desc='This is test case job from Airflow',
iam_role_name='my_test_role',
script_location=some_script,
s3_bucket=some_s3_bucket,
region_name=self.some_aws_region,
)
glue_job_run = glue_job_hook.initialize_job(some_script_arguments)
glue_job_run_state = glue_job_hook.get_job_state(
glue_job_run['JobName'], glue_job_run['JobRunId'])
self.assertEqual(glue_job_run_state,
mock_job_run_state,
msg='Mocks but be equal')
def poke(self, context):
hook = AwsGlueJobHook(aws_conn_id=self.aws_conn_id)
self.log.info("Poking for job run status :for Glue Job %s and ID %s", self.job_name, self.run_id)
job_state = hook.get_job_state(job_name=self.job_name, run_id=self.run_id)
if job_state in self.success_states:
self.log.info("Exiting Job %s Run State: %s", self.run_id, job_state)
return True
elif job_state in self.errored_states:
job_error_message = "Exiting Job " + self.run_id + " Run State: " + job_state
raise AirflowException(job_error_message)
else:
return False