Add StreamInput.readEnum and StreamOutput.writeEnum (#24475)

Implements the common enum serialization/deserialization pattern for enumeration on the StreamInput/StreamOutput.

1 files changed, 30 insertions, 0 deletions

diff --git a/core/src/test/java/org/elasticsearch/common/io/stream/BytesStreamsTests.java b/core/src/test/java/org/elasticsearch/common/io/stream/BytesStreamsTests.java
index a8b065343c..b67000e2b2 100644
--- a/core/src/test/java/org/elasticsearch/common/io/stream/BytesStreamsTests.java
+++ b/core/src/test/java/org/elasticsearch/common/io/stream/BytesStreamsTests.java
@@ -812,4 +812,34 @@ public class BytesStreamsTests extends ESTestCase {
         StreamInput in = new BytesArray(Base64.getDecoder().decode("////////////AQAAAAAAAA==")).streamInput();
         assertEquals(-1, in.readVLong());
     }
+
+    public enum TestEnum {
+        ONE,
+        TWO,
+        THREE
+    }
+
+    public void testEnum() throws IOException {
+        TestEnum value = randomFrom(TestEnum.values());
+        BytesStreamOutput output = new BytesStreamOutput();
+        output.writeEnum(value);
+        StreamInput input = output.bytes().streamInput();
+        assertEquals(value, input.readEnum(TestEnum.class));
+        assertEquals(0, input.available());
+    }
+
+    public void testInvalidEnum() throws IOException {
+        BytesStreamOutput output = new BytesStreamOutput();
+        int randomNumber = randomInt();
+        boolean validEnum = randomNumber >= 0 && randomNumber < TestEnum.values().length;
+        output.writeVInt(randomNumber);
+        StreamInput input = output.bytes().streamInput();
+        if (validEnum) {
+            assertEquals(TestEnum.values()[randomNumber], input.readEnum(TestEnum.class));
+        } else {
+            IOException ex = expectThrows(IOException.class, () -> input.readEnum(TestEnum.class));
+            assertEquals("Unknown TestEnum ordinal [" + randomNumber + "]", ex.getMessage());
+        }
+        assertEquals(0, input.available());
+    }
 }