def test_delete_model():
sagemaker_session = empty_sagemaker_session()
predictor = Predictor(ENDPOINT, sagemaker_session=sagemaker_session)
predictor.delete_model()
expected_call_count = 2
expected_call_args_list = [call("model-1"), call("model-2")]
assert sagemaker_session.delete_model.call_count == expected_call_count
assert sagemaker_session.delete_model.call_args_list == expected_call_args_list
def test_delete_model_fail():
sagemaker_session = empty_sagemaker_session()
sagemaker_session.sagemaker_client.delete_model = Mock(
side_effect=Exception("Could not find model."))
expected_error_message = "One or more models cannot be deleted, please retry."
predictor = Predictor(ENDPOINT, sagemaker_session=sagemaker_session)
with pytest.raises(Exception) as exception:
predictor.delete_model()
assert expected_error_message in str(exception.val)
