From 0b253e5ec663f26464c495f79f46a0af1a5d1f2b Mon Sep 17 00:00:00 2001
From: lxning <23464292+lxning@users.noreply.github.com>
Date: Fri, 30 Apr 2021 20:31:09 -0700
Subject: [PATCH] issue_961 (#969)

* fix prediction string format

* fix test

* fix test
---
 .../com/amazonaws/ml/mms/ModelServerTest.java   | 17 +++--------------
 mms/service.py                                  |  2 +-
 2 files changed, 4 insertions(+), 15 deletions(-)

diff --git a/frontend/server/src/test/java/com/amazonaws/ml/mms/ModelServerTest.java b/frontend/server/src/test/java/com/amazonaws/ml/mms/ModelServerTest.java
index 92639696e..038be0f22 100644
--- a/frontend/server/src/test/java/com/amazonaws/ml/mms/ModelServerTest.java
+++ b/frontend/server/src/test/java/com/amazonaws/ml/mms/ModelServerTest.java
@@ -1219,20 +1219,9 @@ private void testPredictionCustomErrorCode() throws InterruptedException {
         channel.writeAndFlush(req);
         latch.await();
 
-        Assert.assertEquals(httpStatus.code(), 599);
-        channel.close();
-
-        // Unload the model
-        channel = connect(true);
-        httpStatus = null;
-        latch = new CountDownLatch(1);
-        Assert.assertNotNull(channel);
-        req =
-                new DefaultFullHttpRequest(
-                        HttpVersion.HTTP_1_1, HttpMethod.DELETE, "/models/custom-return-code");
-        channel.writeAndFlush(req);
-        latch.await();
-        Assert.assertEquals(httpStatus, HttpResponseStatus.OK);
+        ErrorResponse resp = JsonUtils.GSON.fromJson(result, ErrorResponse.class);
+        Assert.assertEquals(resp.getMessage(), "Some Prediction Error");
+        Assert.assertEquals(resp.getCode(), 599);
     }
 
     private void testErrorBatch() throws InterruptedException {
diff --git a/mms/service.py b/mms/service.py
index 82590c023..1ebe47cc0 100644
--- a/mms/service.py
+++ b/mms/service.py
@@ -138,7 +138,7 @@ def __init__(self, message, error_code=500):
         super(PredictionException, self).__init__(message)
 
     def __str__(self):
-        return "message : error_code".format(message=self.message, error_code=self.error_code)
+        return "{message} : {error_code}".format(message=self.message, error_code=self.error_code)
 
 
 def emit_metrics(metrics):