issue_961 (#969)

* fix prediction string format * fix test * fix test
awslabs · May 1, 2021 · 0b253e5 · 0b253e5
1 parent 46dd188
commit 0b253e5
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 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
@@ -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
@@ -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):