Beispiel #1
0
 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)
Beispiel #2
0
    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,
        )