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"
with pytest.raises(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_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_execute_with_failure(self, mock_processing, mock_client):
sagemaker = SageMakerProcessingOperator(
**self.processing_config_kwargs, config=create_processing_params)
with pytest.raises(AirflowException):
sagemaker.execute(None)