def test_data_wrangler_processor_with_required_parameters(sagemaker_session):
processor = DataWranglerProcessor(
role=ROLE,
data_wrangler_flow_source=DATA_WRANGLER_RECIPE_SOURCE,
instance_count=1,
instance_type="ml.m4.xlarge",
sagemaker_session=sagemaker_session,
)
processor.run()
expected_args = _get_expected_args(processor._current_job_name)
sagemaker_session.process.assert_called_with(**expected_args)
def test_data_wrangler_processor_with_mock_input(sagemaker_session):
processor = DataWranglerProcessor(
role=ROLE,
data_wrangler_flow_source=DATA_WRANGLER_RECIPE_SOURCE,
instance_count=1,
instance_type="ml.m4.xlarge",
sagemaker_session=sagemaker_session,
)
mock_input = ProcessingInput(
source=MOCK_S3_URI,
destination="/opt/ml/processing/mock_input",
input_name="mock_input",
s3_data_type="S3Prefix",
s3_input_mode="File",
s3_data_distribution_type="FullyReplicated",
)
processor.run(inputs=[mock_input])
expected_args = _get_expected_args(processor._current_job_name,
add_mock_input=True)
sagemaker_session.process.assert_called_with(**expected_args)