Octet should accept byte array payload

SeldonIO · Apr 2, 2020 · ea62c09 · ea62c09
1 parent c4952ee
commit ea62c09
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 61 deletions.
diff --git a/engine/src/main/java/io/seldon/engine/api/rest/RestClientController.java b/engine/src/main/java/io/seldon/engine/api/rest/RestClientController.java
@@ -165,7 +165,7 @@ public ResponseEntity<String> predictions_text(RequestEntity<String> requestEnti
           method = RequestMethod.POST,
           consumes = "application/octet-stream",
           produces = "application/json; charset=utf-8")
-  public ResponseEntity<String> predictions_binary(RequestEntity<InputStream> requestEntity) {
+  public ResponseEntity<String> predictions_binary(RequestEntity<byte[]> requestEntity) {
     logger.debug("Received binary predict request");
     Span tracingSpan = null;
     if (tracingProvider.isActive()) {
@@ -174,10 +174,7 @@ public ResponseEntity<String> predictions_binary(RequestEntity<InputStream> requ
       tracer.scopeManager().activate(tracingSpan);
     }
     try {
-      return _predictions(toByteArray(requestEntity.getBody()));
-    } catch (IOException e) {
-      logger.error("Bad request", e);
-      throw new APIException(ApiExceptionType.REQUEST_IO_EXCEPTION, e.getMessage());
+      return _predictions(requestEntity.getBody());
     } finally {
       if (tracingSpan != null) {
         tracingSpan.finish();
@@ -276,6 +273,31 @@ private ResponseEntity<String> _predictions(String json) {
       throw new APIException(ApiExceptionType.ENGINE_INVALID_JSON, json);
     }
 
+    return _predictions(request);
+  }
+
+  /**
+   * It calls the prediction service for the input byte array.
+   *
+   * @param bytes - Input byte array to predict REST api
+   * @return The response for prediction service
+   */
+  private ResponseEntity<String> _predictions(byte[] bytes) {
+    SeldonMessage request = SeldonMessage.newBuilder()
+            .setBinData(ByteString.copyFrom(bytes))
+            .build();
+
+    return _predictions(request);
+  }
+
+  /**
+   * It calls the prediction service for the request. It is the base function for all forms of
+   * request Content-type
+   *
+   * @param request - The SeldonMessage request to the predict REST api
+   * @return The response for prediction service
+   */
+  private ResponseEntity<String> _predictions(SeldonMessage request) {
     try {
       SeldonMessage response = predictionService.predict(request);
       String responseJson = ProtoBufUtils.toJson(response);

diff --git a/engine/src/main/java/io/seldon/engine/config/MessageConvertersConfig.java b/engine/src/main/java/io/seldon/engine/config/MessageConvertersConfig.java
diff --git a/engine/src/test/java/io/seldon/engine/api/rest/TestRestClientController.java b/engine/src/test/java/io/seldon/engine/api/rest/TestRestClientController.java
@@ -347,7 +347,7 @@ public void testPredict_b64img_as_text() throws Exception {
     MvcResult res =
             mvc.perform(
                     MockMvcRequestBuilders.post("/api/v1.0/predictions")
-                            .accept(MediaType.APPLICATION_JSON)
+                            .accept(MediaType.APPLICATION_JSON_UTF8)
                             .content(base64Image)
                             .contentType(MediaType.TEXT_PLAIN))
                     .andReturn();
@@ -364,7 +364,7 @@ public void testPredict_b64img_as_text() throws Exception {
     Assert.assertEquals(base64Image, seldonMessage.getStrData());
     // No Puid specified in request, verify response generated random of correct length
     Assert.assertNotNull(seldonMessage.getMeta().getPuid());
-    Assert.assertTrue(Pattern.matches("[a-z0-7]{26}", seldonMessage.getMeta().getPuid()));
+    Assert.assertTrue(Pattern.matches("[a-z0-9]{26}", seldonMessage.getMeta().getPuid()));
   }
 
   @Test
@@ -373,15 +373,15 @@ public void testPredict_img_as_binary() throws Exception {
     MvcResult res =
             mvc.perform(
                     MockMvcRequestBuilders.post("/api/v1.0/predictions")
-                            .accept(MediaType.APPLICATION_JSON)
+                            .accept(MediaType.APPLICATION_JSON_UTF8)
                             .content(imageBytes)
                             .contentType(MediaType.APPLICATION_OCTET_STREAM))
                     .andReturn();
-    String response = res.getResponse().getContentAsString();
-    System.out.println(response);
+    byte[] response = res.getResponse().getContentAsByteArray();
+    System.out.println(String(response);
     Assert.assertEquals(200, res.getResponse().getStatus());
     SeldonMessage.Builder builder = SeldonMessage.newBuilder();
-    ProtoBufUtils.updateMessageBuilderFromJson(builder, response);
+    ProtoBufUtils.updateMessageBuilderFromJson(builder, new String(response));
     SeldonMessage seldonMessage = builder.build();
     Assert.assertEquals(3, seldonMessage.getMeta().getMetricsCount());
     Assert.assertEquals("COUNTER", seldonMessage.getMeta().getMetrics(0).getType().toString());
@@ -390,6 +390,6 @@ public void testPredict_img_as_binary() throws Exception {
     Assert.assertEquals(imageBytes, seldonMessage.getBinData().toByteArray());
     // No Puid specified in request, verify response generated random of correct length
     Assert.assertNotNull(seldonMessage.getMeta().getPuid());
-    Assert.assertTrue(Pattern.matches("[a-z0-7]{26}", seldonMessage.getMeta().getPuid()));
+    Assert.assertTrue(Pattern.matches("[a-z0-9]{26}", seldonMessage.getMeta().getPuid()));
   }
 }