def test_execute_with_existing_job_fail(self, mock_create_processing_job, mock_list_processing_jobs, mock_client): sagemaker = SageMakerProcessingOperator( **self.processing_config_kwargs, config=create_processing_params) sagemaker.action_if_job_exists = "fail" self.assertRaises(AirflowException, sagemaker.execute, None)
def test_execute(self, mock_processing, mock_client): sagemaker = SageMakerProcessingOperator( **self.processing_config_kwargs, config=create_processing_params) sagemaker.execute(None) mock_processing.assert_called_once_with(create_processing_params, wait_for_completion=False, check_interval=5, max_ingestion_time=None)
def test_action_if_job_exists_validation(self, mock_client): with pytest.raises(AirflowException): SageMakerProcessingOperator( **self.processing_config_kwargs, config=create_processing_params, action_if_job_exists="not_fail_or_increment", )
def test_execute_with_existing_job_increment(self, mock_create_processing_job, mock_list_processing_jobs, mock_client): sagemaker = SageMakerProcessingOperator( **self.processing_config_kwargs, config=create_processing_params) sagemaker.action_if_job_exists = "increment" sagemaker.execute(None) expected_config = create_processing_params.copy() # Expect to see ProcessingJobName suffixed with "-2" because we return one existing job expected_config["ProcessingJobName"] = f"{job_name}-2" mock_create_processing_job.assert_called_once_with( expected_config, wait_for_completion=False, check_interval=5, max_ingestion_time=None, )
def test_action_if_job_exists_validation(self, mock_client): sagemaker = SageMakerProcessingOperator( **self.processing_config_kwargs, config=create_processing_params) self.assertRaises(AirflowException, sagemaker.__init__, action_if_job_exists="not_fail_or_increment")
def test_execute_with_failure(self, mock_processing, mock_client): sagemaker = SageMakerProcessingOperator( **self.processing_config_kwargs, config=create_processing_params) self.assertRaises(AirflowException, sagemaker.execute, None)
def test_integer_fields_are_set(self, config, expected_fields): sagemaker = SageMakerProcessingOperator(** self.processing_config_kwargs, config=config) assert sagemaker.integer_fields == expected_fields