class TestSageMakerModelOperator(unittest.TestCase):
def setUp(self):
self.sagemaker = SageMakerModelOperator(
task_id='test_sagemaker_operator',
aws_conn_id='sagemaker_test_id',
config=create_model_params)
@mock.patch.object(SageMakerHook, 'get_conn')
@mock.patch.object(SageMakerHook, 'create_model')
def test_execute(self, mock_model, mock_client):
mock_model.return_value = {
'ModelArn': 'testarn',
'ResponseMetadata': {
'HTTPStatusCode': 200
}
}
self.sagemaker.execute(None)
mock_model.assert_called_once_with(create_model_params)
@mock.patch.object(SageMakerHook, 'get_conn')
@mock.patch.object(SageMakerHook, 'create_model')
def test_execute_with_failure(self, mock_model, mock_client):
mock_model.return_value = {
'ModelArn': 'testarn',
'ResponseMetadata': {
'HTTPStatusCode': 404
}
}
self.assertRaises(AirflowException, self.sagemaker.execute, None)
def setUp(self):
self.sagemaker = SageMakerModelOperator(
task_id='test_sagemaker_operator', aws_conn_id='sagemaker_test_id', config=create_model_params
)
