Example #1
0
    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')
Example #2
0
 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