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)