Fix netty tests

core/server/worker/src/test/java/alluxio/worker/netty/DataServerBlockReadHandlerTest.java

@@ -14,6 +14,7 @@
 import alluxio.EmbeddedNoExceptionChannel;
 import alluxio.network.protocol.RPCProtoMessage;
 import alluxio.proto.dataserver.Protocol;
+import alluxio.proto.status.Status.PStatus;
 import alluxio.util.proto.ProtoMessage;
 import alluxio.worker.block.BlockWorker;
 import alluxio.worker.block.io.BlockReader;
@@ -80,7 +81,7 @@ public void readFailure() throws Exception {
     mBlockReader.close();
     mChannelNoException.writeInbound(buildReadRequest(0, fileSize));
     Object response = waitForOneResponse(mChannelNoException);
-    checkReadResponse(response, Protocol.Status.Code.INTERNAL);
+    checkReadResponse(response, PStatus.FAILED_PRECONDITION);
   }
 
   @Override

core/server/worker/src/test/java/alluxio/worker/netty/DataServerBlockWriteHandlerTest.java

@@ -16,6 +16,7 @@
 import alluxio.network.protocol.databuffer.DataBuffer;
 import alluxio.network.protocol.databuffer.DataNettyBufferV2;
 import alluxio.proto.dataserver.Protocol;
+import alluxio.proto.status.Status.PStatus;
 import alluxio.util.io.BufferUtils;
 import alluxio.util.proto.ProtoMessage;
 import alluxio.worker.block.BlockWorker;
@@ -71,7 +72,7 @@ public void writeFailure() throws Exception {
     mChannelNoException.writeInbound(buildWriteRequest(PACKET_SIZE, PACKET_SIZE));
     Object writeResponse = waitForResponse(mChannelNoException);
     Assert.assertTrue(writeResponse instanceof RPCProtoMessage);
-    checkWriteResponse(writeResponse, Protocol.Status.Code.INTERNAL);
+    checkWriteResponse(writeResponse, PStatus.FAILED_PRECONDITION);
   }
 
   @Override

core/server/worker/src/test/java/alluxio/worker/netty/DataServerReadHandlerTest.java

@@ -14,12 +14,12 @@
 import alluxio.Configuration;
 import alluxio.Constants;
 import alluxio.PropertyKey;
-import alluxio.exception.status.Status;
 import alluxio.network.protocol.RPCProtoMessage;
 import alluxio.network.protocol.databuffer.DataBuffer;
 import alluxio.network.protocol.databuffer.DataFileChannel;
 import alluxio.network.protocol.databuffer.DataNettyBufferV2;
 import alluxio.proto.dataserver.Protocol;
+import alluxio.proto.status.Status.PStatus;
 import alluxio.util.CommonUtils;
 import alluxio.util.WaitForOptions;
 import alluxio.util.io.BufferUtils;
@@ -98,7 +98,7 @@ public void readEmptyFile() throws Exception {
     populateInputFile(0, 0, 0);
     mChannelNoException.writeInbound(buildReadRequest(0, 0));
     Object response = waitForOneResponse(mChannelNoException);
-    checkReadResponse(response, Protocol.Status.Code.INVALID_ARGUMENT);
+    checkReadResponse(response, PStatus.INVALID_ARGUMENT);
   }
 
   /**
@@ -124,7 +124,7 @@ public void cancelRequest() throws Exception {
       // There is small chance that we can still receive an OK response here because it is too
       // fast to read all the data. If that ever happens, either increase the file size or allow it
       // to be OK here.
-      DataBuffer buffer = checkReadResponse(response, Protocol.Status.Code.CANCELED);
+      DataBuffer buffer = checkReadResponse(response, PStatus.CANCELED);
       if (buffer == null) {
         eof = true;
         break;
@@ -183,7 +183,7 @@ protected void checkAllReadResponses(EmbeddedChannel channel, long checksumExpec
         Assert.fail();
         break;
       }
-      DataBuffer buffer = checkReadResponse(readResponse, Protocol.Status.Code.OK);
+      DataBuffer buffer = checkReadResponse(readResponse, PStatus.OK);
       eof = buffer == null;
       if (buffer != null) {
         if (buffer instanceof DataNettyBufferV2) {
@@ -214,18 +214,17 @@ protected void checkAllReadResponses(EmbeddedChannel channel, long checksumExpec
    * @param statusExpected the expected error code
    * @return the data buffer extracted from the read response
    */
-  protected DataBuffer checkReadResponse(Object readResponse, Status statusExpected) {
+  protected DataBuffer checkReadResponse(Object readResponse, PStatus statusExpected) {
     Assert.assertTrue(readResponse instanceof RPCProtoMessage);
 
     ProtoMessage response = ((RPCProtoMessage) readResponse).getMessage();
     Assert.assertTrue(response.getType() == ProtoMessage.Type.RESPONSE);
     DataBuffer buffer = ((RPCProtoMessage) readResponse).getPayloadDataBuffer();
     if (buffer != null) {
-      Assert.assertEquals(Status.OK,
-          response.<Protocol.Response>getMessage().getStatus().getCode());
+      Assert.assertEquals(PStatus.OK, response.<Protocol.Response>getMessage().getStatus());
     } else {
       Assert.assertEquals(statusExpected,
-          response.<Protocol.Response>getMessage().getStatus().getCode());
+          response.<Protocol.Response>getMessage().getStatus());
     }
     return buffer;
   }

core/server/worker/src/test/java/alluxio/worker/netty/DataServerUFSFileReadHandlerTest.java

@@ -13,8 +13,9 @@
 
 import alluxio.EmbeddedNoExceptionChannel;
 import alluxio.network.protocol.RPCProtoMessage;
-import alluxio.util.proto.ProtoMessage;
 import alluxio.proto.dataserver.Protocol;
+import alluxio.proto.status.Status.PStatus;
+import alluxio.util.proto.ProtoMessage;
 import alluxio.worker.file.FileSystemWorker;
 
 import io.netty.channel.embedded.EmbeddedChannel;
@@ -58,7 +59,7 @@ public void readFailure() throws Exception {
     mInputStream.close();
     mChannelNoException.writeInbound(buildReadRequest(0, fileSize));
     Object response = waitForOneResponse(mChannelNoException);
-    checkReadResponse(response, Protocol.Status.Code.INTERNAL);
+    checkReadResponse(response, PStatus.UNAVAILABLE);
   }
 
   @Override

core/server/worker/src/test/java/alluxio/worker/netty/DataServerUFSFileWriteHandlerTest.java

@@ -16,6 +16,7 @@
 import alluxio.network.protocol.databuffer.DataBuffer;
 import alluxio.network.protocol.databuffer.DataNettyBufferV2;
 import alluxio.proto.dataserver.Protocol;
+import alluxio.proto.status.Status.PStatus;
 import alluxio.underfs.UnderFileSystem;
 import alluxio.underfs.options.CreateOptions;
 import alluxio.util.io.BufferUtils;
@@ -78,7 +79,7 @@ public void writeFailure() throws Exception {
     mOutputStream.close();
     mChannelNoException.writeInbound(buildWriteRequest(PACKET_SIZE, PACKET_SIZE));
     Object writeResponse = waitForResponse(mChannelNoException);
-    checkWriteResponse(writeResponse, Protocol.Status.Code.INTERNAL);
+    checkWriteResponse(writeResponse, PStatus.UNAVAILABLE);
   }
 
   @Override

core/server/worker/src/test/java/alluxio/worker/netty/DataServerWriteHandlerTest.java

@@ -12,9 +12,9 @@
 package alluxio.worker.netty;
 
 import alluxio.Constants;
-import alluxio.exception.status.Status;
 import alluxio.network.protocol.RPCProtoMessage;
 import alluxio.proto.dataserver.Protocol;
+import alluxio.proto.status.Status.PStatus;
 import alluxio.util.CommonUtils;
 import alluxio.util.WaitForOptions;
 import alluxio.util.io.BufferUtils;
@@ -56,7 +56,7 @@ public void writeEmptyFile() throws Exception {
     mChannel.writeInbound(buildWriteRequest(0, 0));
 
     Object writeResponse = waitForResponse(mChannel);
-    checkWriteResponse(writeResponse, Status.OK);
+    checkWriteResponse(writeResponse, PStatus.OK);
   }
 
   /**
@@ -73,7 +73,7 @@ public void writeNonEmptyFile() throws Exception {
     mChannel.writeInbound(buildWriteRequest(len, 0));
 
     Object writeResponse = waitForResponse(mChannel);
-    checkWriteResponse(writeResponse, Status.OK);
+    checkWriteResponse(writeResponse, PStatus.OK);
     checkFileContent(len);
   }
 
@@ -91,7 +91,7 @@ public void cancel() throws Exception {
     mChannel.writeInbound(buildWriteRequest(len, -1));
 
     Object writeResponse = waitForResponse(mChannel);
-    checkWriteResponse(writeResponse, Status.CANCELED);
+    checkWriteResponse(writeResponse, PStatus.CANCELED);
     // Our current implementation does not really abort the file when the write is cancelled.
     // The client issues another request to block worker to abort it.
     checkFileContent(len);
@@ -106,7 +106,7 @@ public void writeInvalidOffset() throws Exception {
     mChannelNoException.writeInbound(buildWriteRequest(PACKET_SIZE + 1, PACKET_SIZE));
     Object writeResponse = waitForResponse(mChannelNoException);
     Assert.assertTrue(writeResponse instanceof RPCProtoMessage);
-    checkWriteResponse(writeResponse, Status.INVALID_ARGUMENT);
+    checkWriteResponse(writeResponse, PStatus.INVALID_ARGUMENT);
   }
 
   /**
@@ -115,13 +115,12 @@ public void writeInvalidOffset() throws Exception {
    * @param writeResponse the write response
    * @param statusExpected the expected status code
    */
-  protected void checkWriteResponse(Object writeResponse, Status statusExpected) {
+  protected void checkWriteResponse(Object writeResponse, PStatus statusExpected) {
     Assert.assertTrue(writeResponse instanceof RPCProtoMessage);
 
     ProtoMessage response = ((RPCProtoMessage) writeResponse).getMessage();
     Assert.assertTrue(response.getType() == ProtoMessage.Type.RESPONSE);
-    Assert.assertEquals(Status.toProto(statusExpected),
-        response.<Protocol.Response>getMessage().getException().getStatus());
+    Assert.assertEquals(statusExpected, response.<Protocol.Response>getMessage().getStatus());
   }
 
   /**