fix wired import format

SeldonIO · Oct 29, 2019 · a6a9f2f · a6a9f2f
1 parent a2fe911
commit a6a9f2f
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 13 deletions.
diff --git a/engine/src/main/java/io/seldon/engine/service/InternalPredictionService.java b/engine/src/main/java/io/seldon/engine/service/InternalPredictionService.java
@@ -36,11 +36,11 @@
 import io.seldon.protos.ModelGrpc.ModelBlockingStub;
 import io.seldon.protos.OutputTransformerGrpc;
 import io.seldon.protos.OutputTransformerGrpc.OutputTransformerBlockingStub;
-import io.seldon.protos.PredictionProtos;
 import io.seldon.protos.PredictionProtos.Feedback;
 import io.seldon.protos.PredictionProtos.SeldonMessage;
 import io.seldon.protos.PredictionProtos.SeldonMessage.DataOneofCase;
 import io.seldon.protos.PredictionProtos.SeldonMessageList;
+import io.seldon.protos.PredictionProtos.Status;
 import io.seldon.protos.RouterGrpc;
 import io.seldon.protos.RouterGrpc.RouterBlockingStub;
 import io.seldon.protos.TransformerGrpc;
@@ -502,7 +502,7 @@ private SeldonMessage queryREST(
             logger.error(
                 "Couldn't retrieve prediction from external prediction server -- bad http return code: "
                     + httpResponse.getStatusCode());
-            PredictionProtos.Status seldonMessageStatus = seldonMessage.getStatus();
+            Status seldonMessageStatus = seldonMessage.getStatus();
             if (seldonMessageStatus == null) {
                 throw new APIException(
                         APIException.ApiExceptionType.ENGINE_MICROSERVICE_ERROR,

diff --git a/engine/src/test/java/io/seldon/engine/ExceptionControllerAdviceTest.java b/engine/src/test/java/io/seldon/engine/ExceptionControllerAdviceTest.java
@@ -2,7 +2,9 @@
 
 import com.google.protobuf.util.JsonFormat;
 import io.seldon.engine.exception.APIException;
-import io.seldon.protos.PredictionProtos;
+import io.seldon.engine.exception.APIException.ApiExceptionType;
+import io.seldon.protos.PredictionProtos.SeldonMessage;
+import io.seldon.protos.PredictionProtos.Status;
 import org.junit.Assert;
 import org.junit.Test;
 import org.springframework.http.ResponseEntity;
@@ -11,33 +13,32 @@ public class ExceptionControllerAdviceTest {
     @Test
     public void testApiExceptionType() throws Exception {
         ResponseEntity<String> responseEntity = new io.seldon.engine.ExceptionControllerAdvice()
-                .handleUnauthorizedException(new APIException(APIException
-                        .ApiExceptionType.ENGINE_MICROSERVICE_ERROR, "info"));
-        validateSeldonMessage(responseEntity, APIException
-                .ApiExceptionType.ENGINE_MICROSERVICE_ERROR);
+                .handleUnauthorizedException(new APIException(
+                        ApiExceptionType.ENGINE_MICROSERVICE_ERROR, "info"));
+        validateSeldonMessage(responseEntity, ApiExceptionType.ENGINE_MICROSERVICE_ERROR);
     }
 
     @Test
     public void testCustomizedExceptionType() throws Exception {
-        APIException.ApiExceptionType exceptionType =
-                APIException.ApiExceptionType.CUSTOMIZED_EXCEPTION;
+        ApiExceptionType exceptionType =
+                ApiExceptionType.CUSTOMIZED_EXCEPTION;
         exceptionType.setMessage("exception msg in test");
         ResponseEntity<String> responseEntity = new io.seldon.engine.ExceptionControllerAdvice()
                 .handleUnauthorizedException(new APIException(exceptionType, "info"));
         validateSeldonMessage(responseEntity, exceptionType);
     }
 
     private void validateSeldonMessage(
-            ResponseEntity<String> httpResponse, APIException.ApiExceptionType exceptionType) throws Exception {
+            ResponseEntity<String> httpResponse, ApiExceptionType exceptionType) throws Exception {
         String response = httpResponse.getBody();
-        PredictionProtos.SeldonMessage.Builder builder = PredictionProtos.SeldonMessage.newBuilder();
+        SeldonMessage.Builder builder = SeldonMessage.newBuilder();
         JsonFormat.parser().ignoringUnknownFields().merge(response, builder);
-        PredictionProtos.SeldonMessage seldonMessage = builder.build();
+        SeldonMessage seldonMessage = builder.build();
 
         Assert.assertEquals(exceptionType.getHttpCode(), httpResponse.getStatusCodeValue());
         Assert.assertEquals(exceptionType.getId(), seldonMessage.getStatus().getCode());
         Assert.assertEquals(exceptionType.getMessage(), seldonMessage.getStatus().getReason());
         Assert.assertEquals("info", seldonMessage.getStatus().getInfo());
-        Assert.assertEquals(PredictionProtos.Status.StatusFlag.FAILURE, seldonMessage.getStatus().getStatus());
+        Assert.assertEquals(Status.StatusFlag.FAILURE, seldonMessage.getStatus().getStatus());
     }
 }