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)