Test improvements

ambry-router/src/test/java/com.github.ambry.router/GetBlobInfoOperationTest.java

@@ -57,7 +57,7 @@ public class GetBlobInfoOperationTest {
 
   private int requestParallelism = 2;
   private int successTarget = 1;
-  private final RouterConfig routerConfig;
+  private RouterConfig routerConfig;
   private final MockClusterMap mockClusterMap;
   private final MockServerLayout mockServerLayout;
   private final int replicasCount;
@@ -89,7 +89,7 @@ public void registerRequestToSend(GetOperation getOperation, RequestInfo request
 
   public GetBlobInfoOperationTest()
       throws Exception {
-    VerifiableProperties vprops = getNonBlockingRouterProperties();
+    VerifiableProperties vprops = new VerifiableProperties(getNonBlockingRouterProperties());
     routerConfig = new RouterConfig(vprops);
     mockClusterMap = new MockClusterMap();
     mockServerLayout = new MockServerLayout(mockClusterMap);
@@ -338,6 +338,28 @@ private void testErrorPrecedence(ServerErrorCode[] serverErrorCodesInOrder, Rout
   @Test
   public void testSuccessInThePresenceOfVariousErrors()
       throws Exception {
+    // The put for the blob being requested happened.
+    String dcWherePutHappened = routerConfig.routerDatacenterName;
+
+    // test requests coming in from local dc as well as cross-colo.
+    Properties props = getNonBlockingRouterProperties();
+    props.setProperty("router.datacenter.name", "DC1");
+    routerConfig = new RouterConfig(new VerifiableProperties(props));
+    testVariousErrors(dcWherePutHappened);
+
+    props = getNonBlockingRouterProperties();
+    props.setProperty("router.datacenter.name", "DC2");
+    routerConfig = new RouterConfig(new VerifiableProperties(props));
+    testVariousErrors(dcWherePutHappened);
+
+    props = getNonBlockingRouterProperties();
+    props.setProperty("router.datacenter.name", "DC3");
+    routerConfig = new RouterConfig(new VerifiableProperties(props));
+    testVariousErrors(dcWherePutHappened);
+  }
+
+  private void testVariousErrors(String dcWherePutHappened)
+      throws Exception {
     GetBlobInfoOperation op =
         new GetBlobInfoOperation(routerConfig, mockClusterMap, responseHandler, blobIdStr, operationFuture, null, time);
     ArrayList<RequestInfo> requestListToFill = new ArrayList<>();
@@ -355,11 +377,10 @@ public void testSuccessInThePresenceOfVariousErrors()
     mockServers.get(7).setServerErrorForAllRequests(ServerErrorCode.Disk_Unavailable);
     mockServers.get(8).setServerErrorForAllRequests(ServerErrorCode.Unknown_Error);
 
-    // set the status of one of the servers in the local datacenter to success (we depend on an actual put of the blob
-    // in order to get back the BlobInfo, and the put would have gone only to the local datacenter).
+    // clear the hard error in one of the servers in the datacenter where the put happened.
     for (int i = 0; i < mockServers.size(); i++) {
       MockServer mockServer = mockServers.get(i);
-      if (mockServer.getDataCenter().equals(routerConfig.routerDatacenterName)) {
+      if (mockServer.getDataCenter().equals(dcWherePutHappened)) {
         mockServer.setServerErrorForAllRequests(ServerErrorCode.No_Error);
         break;
       }
@@ -417,13 +438,13 @@ private void assertSuccess(GetBlobInfoOperation op) {
    * Get the properties for the {@link NonBlockingRouter}.
    * @return the constructed properties.
    */
-  private VerifiableProperties getNonBlockingRouterProperties() {
+  private Properties getNonBlockingRouterProperties() {
     Properties properties = new Properties();
     properties.setProperty("router.hostname", "localhost");
     properties.setProperty("router.datacenter.name", "DC1");
     properties.setProperty("router.get.request.parallelism", Integer.toString(requestParallelism));
     properties.setProperty("router.get.success.target", Integer.toString(successTarget));
-    return new VerifiableProperties(properties);
+    return properties;
   }
 }
 

ambry-router/src/test/java/com.github.ambry.router/MockServer.java

@@ -39,6 +39,7 @@
 import java.io.DataInputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.nio.channels.Channels;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.LinkedList;
@@ -66,12 +67,12 @@ class MockServer {
 
   /**
    * Take in a request in the form of {@link Send} and return a response in the form of a
-   * {@link MockBoundedByteBufferReceive}.
+   * {@link BoundedByteBufferReceive}.
    * @param send the request.
    * @return the response.
    * @throws IOException if there was an error in interpreting the request.
    */
-  public MockBoundedByteBufferReceive send(Send send)
+  public BoundedByteBufferReceive send(Send send)
       throws IOException {
     if (!shouldRespond) {
       return null;
@@ -97,10 +98,9 @@ public MockBoundedByteBufferReceive send(Send send)
     response.writeTo(channel);
     ByteBuffer payload = channel.getBuffer();
     payload.flip();
-    // read the size off. the size is used by the client to allocate the buffer and then interpret the response.
-    // MockServer abstracts it away at a level above that and returns the "allocated buffer".
-    payload.getLong();
-    return new MockBoundedByteBufferReceive(payload);
+    BoundedByteBufferReceive boundedByteBufferReceive = new BoundedByteBufferReceive();
+    boundedByteBufferReceive.readFrom(Channels.newChannel(new ByteBufferInputStream(payload)));
+    return boundedByteBufferReceive;
   }
 
   PutResponse makePutResponse(PutRequest putRequest, ServerErrorCode putError)
@@ -194,7 +194,8 @@ GetResponse makeGetResponse(GetRequest getRequest, ServerErrorCode getError)
       getResponse = new GetResponse(getRequest.getCorrelationId(), getRequest.getClientId(), partitionResponseInfoList,
           responseSend, serverError);
     } else {
-      getResponse = new GetResponse(getRequest.getCorrelationId(), getRequest.getClientId(), serverError);
+      getResponse = new GetResponse(getRequest.getCorrelationId(), getRequest.getClientId(),
+          new ArrayList<PartitionResponseInfo>(), new ByteBufferSend(ByteBuffer.allocate(0)), serverError);
     }
     return getResponse;
   }
@@ -289,7 +290,8 @@ public ServerErrorCode getErrorFromBlobIdStr(String blobIdString) {
   }
 
   /**
-   * Set the mapping relationship between a {@code blobIdString} and the {@link ServerErrorCode} this server should return.
+   * Set the mapping relationship between a {@code blobIdString} and the {@link ServerErrorCode} this server should
+   * return.
    * @param blobIdString The key in this mapping relation.
    * @param code The {@link ServerErrorCode} for the {@code blobIdString}.
    */
@@ -298,28 +300,3 @@ public void setBlobIdToServerErrorCode(String blobIdString, ServerErrorCode code
   }
 }
 
-/**
- * A mock implementation of {@link BoundedByteBufferReceive} that constructs a buffer with the passed in correlation
- * id and returns that buffer as part of {@link #getPayload()}.
- */
-class MockBoundedByteBufferReceive extends BoundedByteBufferReceive {
-  private final ByteBuffer buf;
-
-  /**
-   * Construct a MockBoundedByteBufferReceive with the given correlation id.
-   * @param buf the ByteBuffer that is the payload of this object.
-   */
-  public MockBoundedByteBufferReceive(ByteBuffer buf) {
-    this.buf = buf;
-  }
-
-  /**
-   * Return the buffer associated with this object.
-   * @return the buffer associated with this object.
-   */
-  @Override
-  public ByteBuffer getPayload() {
-    return buf;
-  }
-}
-