Add unit tests for ServerAdminAction & GrpcUtilsTest

internal/venice-common/src/main/java/com/linkedin/venice/meta/ServerAdminAction.java

@@ -1,8 +1,19 @@
 package com.linkedin.venice.meta;
 
+import java.util.HashMap;
+import java.util.Map;
+
+
 public enum ServerAdminAction {
   DUMP_INGESTION_STATE(0), DUMP_SERVER_CONFIGS(1);
 
+  private static final Map<Integer, ServerAdminAction> ADMIN_ACTION_MAP = new HashMap<>(2);
+
+  static {
+    for (ServerAdminAction action: values()) {
+      ADMIN_ACTION_MAP.put(action.getValue(), action);
+    }
+  }
   private final int value;
 
   ServerAdminAction(int value) {
@@ -13,13 +24,11 @@ public int getValue() {
     return this.value;
   }
 
-  // get the enum value from the integer value
   public static ServerAdminAction fromValue(int value) {
-    for (ServerAdminAction action: ServerAdminAction.values()) {
-      if (action.getValue() == value) {
-        return action;
-      }
+    ServerAdminAction action = ADMIN_ACTION_MAP.get(value);
+    if (action == null) {
+      throw new IllegalArgumentException("Unknown server admin action value: " + value);
     }
-    throw new IllegalArgumentException("Invalid value for ServerAdminAction: " + value);
+    return action;
   }
 }

internal/venice-common/src/test/java/com/linkedin/venice/grpc/GrpcUtilsTest.java

@@ -1,12 +1,19 @@
 package com.linkedin.venice.grpc;
 
-import static org.mockito.Mockito.*;
-import static org.testng.Assert.*;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
 
+import com.google.protobuf.ByteString;
 import com.linkedin.venice.security.SSLFactory;
 import com.linkedin.venice.utils.SslUtils;
 import io.grpc.Status;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.CompositeByteBuf;
+import io.netty.buffer.Unpooled;
 import io.netty.handler.codec.http.HttpResponseStatus;
+import java.nio.ByteBuffer;
 import javax.net.ssl.KeyManager;
 import javax.net.ssl.TrustManager;
 import org.testng.annotations.BeforeTest;
@@ -71,4 +78,98 @@ public void testHttpResponseStatusToGrpcStatus() {
         grpcStatus.getDescription(),
         "Mismatch in error description for the mapped grpc status");
   }
+
+  @Test
+  public void testToByteStringWithArrayBacking() {
+    byte[] byteArray = new byte[] { 1, 2, 3, 4, 5 };
+    ByteBuf byteBuf = Unpooled.wrappedBuffer(byteArray);
+    assertTrue(byteBuf.hasArray());
+
+    ByteString byteString = GrpcUtils.toByteString(byteBuf);
+    assertEquals(byteString.toByteArray(), byteArray);
+  }
+
+  @Test
+  public void testToByteStringWithArrayBackingAndReaderIndex() {
+    byte[] byteArray = new byte[] { 1, 2, 3, 4, 5 };
+    ByteBuf byteBuf = Unpooled.wrappedBuffer(byteArray);
+    byteBuf.readerIndex(2); // Set reader index to 2
+
+    assertTrue(byteBuf.hasArray());
+    ByteString byteString = GrpcUtils.toByteString(byteBuf);
+
+    byte[] expectedArray = new byte[] { 3, 4, 5 }; // Expected array after reader index
+    assertEquals(byteString.toByteArray(), expectedArray);
+  }
+
+  @Test
+  public void testToByteStringWithEmptyByteBuf() {
+    ByteBuf byteBuf = Unpooled.EMPTY_BUFFER;
+
+    ByteString byteString = GrpcUtils.toByteString(byteBuf);
+
+    assertTrue(byteString.isEmpty());
+  }
+
+  @Test
+  public void testToByteStringWithNonArrayBacking() {
+    // Create a direct ByteBuf with an initial capacity of 16 bytes
+    ByteBuf byteBuf = Unpooled.directBuffer(16);
+    // Write some data into the ByteBuf
+    byteBuf.writeBytes(new byte[] { 1, 2, 3, 4, 5 });
+    assertFalse(byteBuf.hasArray());
+
+    ByteString byteString = GrpcUtils.toByteString(byteBuf);
+
+    byte[] expectedArray = new byte[] { 1, 2, 3, 4, 5 };
+    assertEquals(byteString.toByteArray(), expectedArray);
+  }
+
+  @Test
+  public void testToByteStringWithDifferentOffsetsAndLengths() {
+    byte[] byteArray = new byte[] { 10, 20, 30, 40, 50 };
+    ByteBuf byteBuf = Unpooled.wrappedBuffer(byteArray);
+    byteBuf.readerIndex(1); // Start from index 1
+    byteBuf.writerIndex(4); // End at index 4
+
+    ByteString byteString = GrpcUtils.toByteString(byteBuf);
+
+    byte[] expectedArray = new byte[] { 20, 30, 40 }; // Expected byte array
+    assertEquals(byteString.toByteArray(), expectedArray);
+  }
+
+  @Test
+  public void testToByteStringWithNioBuffer() {
+    byte[] byteArray = new byte[] { 100, (byte) 200, (byte) 255 };
+    ByteBuf byteBuf = Unpooled.wrappedBuffer(byteArray);
+
+    ByteBuffer nioBuffer = byteBuf.nioBuffer(); // Get NIO buffer
+    ByteString byteString = GrpcUtils.toByteString(Unpooled.wrappedBuffer(nioBuffer));
+
+    assertEquals(byteString.toByteArray(), byteArray);
+  }
+
+  @Test
+  public void testToByteStringWithCompositeByteBuf() {
+    // Create some ByteBufs
+    ByteBuf buffer1 = Unpooled.wrappedBuffer(new byte[] { 10, 20 });
+    ByteBuf buffer2 = Unpooled.wrappedBuffer(new byte[] { 30, 40, 50 });
+
+    // Create a CompositeByteBuf that combines the above ByteBufs
+    CompositeByteBuf compositeByteBuf = Unpooled.compositeBuffer();
+    compositeByteBuf.addComponents(true, buffer1, buffer2);
+
+    // The expected ByteString combining all bytes from buffer1 and buffer2
+    byte[] expectedBytes = new byte[] { 10, 20, 30, 40, 50 };
+
+    // Verify that the CompositeByteBuf has the combined array
+    byte[] combinedArray = new byte[5];
+    compositeByteBuf.getBytes(0, combinedArray);
+    assertEquals(combinedArray, expectedBytes);
+
+    // Convert CompositeByteBuf to ByteString using the method
+    ByteString byteString = GrpcUtils.toByteString(compositeByteBuf);
+
+    assertEquals(byteString.toByteArray(), expectedBytes);
+  }
 }

internal/venice-common/src/test/java/com/linkedin/venice/meta/ServerAdminActionTest.java

@@ -0,0 +1,20 @@
+package com.linkedin.venice.meta;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.expectThrows;
+
+import org.testng.annotations.Test;
+
+
+public class ServerAdminActionTest {
+  @Test
+  public void testFromValue() {
+    // Test all valid values
+    assertEquals(ServerAdminAction.fromValue(0), ServerAdminAction.DUMP_INGESTION_STATE);
+    assertEquals(ServerAdminAction.fromValue(1), ServerAdminAction.DUMP_SERVER_CONFIGS);
+
+    // Test invalid values
+    expectThrows(IllegalArgumentException.class, () -> ServerAdminAction.fromValue(-1));
+    expectThrows(IllegalArgumentException.class, () -> ServerAdminAction.fromValue(2));
+  }
+}

internal/venice-common/src/test/java/com/linkedin/venice/utils/NettyUtilsTest.java

@@ -0,0 +1,31 @@
+package com.linkedin.venice.utils;
+
+import static com.linkedin.venice.HttpConstants.VENICE_RETRY;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+
+import io.netty.handler.codec.http.DefaultHttpRequest;
+import io.netty.handler.codec.http.HttpMethod;
+import io.netty.handler.codec.http.HttpRequest;
+import io.netty.handler.codec.http.HttpVersion;
+import org.testng.annotations.Test;
+
+
+public class NettyUtilsTest {
+  @Test
+  public void testContainRetryHeader() {
+    // 1. Header is present
+    HttpRequest httpRequest = new DefaultHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/test");
+    httpRequest.headers().set(VENICE_RETRY, "1");
+    assertTrue(NettyUtils.containRetryHeader(httpRequest), "Request should contain the retry header");
+
+    // 2. Header is present with value "true"
+    httpRequest.headers().clear();
+    httpRequest.headers().set(VENICE_RETRY, "true");
+    assertTrue(NettyUtils.containRetryHeader(httpRequest), "Request should contain the retry header");
+
+    // 3. Header is not present
+    httpRequest.headers().clear();
+    assertFalse(NettyUtils.containRetryHeader(httpRequest), "Request should not contain the retry header");
+  }
+}