From d2dac8427e0fb8bdb50e98907ac1a6ffbb1bf929 Mon Sep 17 00:00:00 2001
From: Isabel Drost-Fromm <isabel.drostfromm@elasticsearch.com>
Date: Wed, 20 Jan 2016 15:01:54 +0100
Subject: [PATCH] Make distance unit writable.

Relates to #15178
---
 .../common/unit/DistanceUnit.java             | 23 ++++++++++++++++++-
 .../common/unit/DistanceUnitTests.java        | 21 +++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/core/src/main/java/org/elasticsearch/common/unit/DistanceUnit.java b/core/src/main/java/org/elasticsearch/common/unit/DistanceUnit.java
index feebd93c5ab..d0e91646c0f 100644
--- a/core/src/main/java/org/elasticsearch/common/unit/DistanceUnit.java
+++ b/core/src/main/java/org/elasticsearch/common/unit/DistanceUnit.java
@@ -22,6 +22,7 @@ package org.elasticsearch.common.unit;
 import org.elasticsearch.common.geo.GeoUtils;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
 
 import java.io.IOException;
 
@@ -32,7 +33,7 @@ import java.io.IOException;
  * the earth ellipsoid defined in {@link GeoUtils}. The default unit used within
  * this project is <code>METERS</code> which is defined by <code>DEFAULT</code>
  */
-public enum DistanceUnit {
+public enum DistanceUnit implements Writeable<DistanceUnit> {
     INCH(0.0254, "in", "inch"),
     YARD(0.9144, "yd", "yards"),
     FEET(0.3048, "ft", "feet"),
@@ -322,4 +323,24 @@ public enum DistanceUnit {
             return new Distance(Double.parseDouble(distance), defaultUnit);
         }
     }
+
+    private static final DistanceUnit PROTOTYPE = DEFAULT;
+
+    @Override
+    public DistanceUnit readFrom(StreamInput in) throws IOException {
+        int ordinal = in.readVInt();
+        if (ordinal < 0 || ordinal >= values().length) {
+            throw new IOException("Unknown DistanceUnit ordinal [" + ordinal + "]");
+        }
+        return values()[ordinal];
+    }
+
+    public static DistanceUnit readUnitFrom(StreamInput in) throws IOException {
+        return PROTOTYPE.readFrom(in);
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeVInt(this.ordinal());
+    }
 }
diff --git a/core/src/test/java/org/elasticsearch/common/unit/DistanceUnitTests.java b/core/src/test/java/org/elasticsearch/common/unit/DistanceUnitTests.java
index 25c3a136271..5d7bbb3ca18 100644
--- a/core/src/test/java/org/elasticsearch/common/unit/DistanceUnitTests.java
+++ b/core/src/test/java/org/elasticsearch/common/unit/DistanceUnitTests.java
@@ -19,6 +19,10 @@
 
 package org.elasticsearch.common.unit;
 
+import com.carrotsearch.randomizedtesting.generators.RandomStrings;
+
+import org.elasticsearch.common.io.stream.BytesStreamOutput;
+import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.test.ESTestCase;
 
 import static org.hamcrest.Matchers.closeTo;
@@ -73,4 +77,21 @@ public class DistanceUnitTests extends ESTestCase {
         assertEquals(7, DistanceUnit.MILES.ordinal());
         assertEquals(8, DistanceUnit.METERS.ordinal());
     }
+
+    public void testReadWrite() throws Exception {
+        for (DistanceUnit unit : DistanceUnit.values()) {
+          try (BytesStreamOutput out = new BytesStreamOutput()) {
+              unit.writeTo(out);
+              try (StreamInput in = StreamInput.wrap(out.bytes())) {
+                  assertThat("Roundtrip serialisation failed.", DistanceUnit.readDistanceUnit(in), equalTo(unit));
+              }
+          }
+        }
+    }
+
+    public void testFromString() {
+        for (DistanceUnit unit : DistanceUnit.values()) {
+            assertThat("Roundtrip string parsing failed.", DistanceUnit.fromString(unit.toString()), equalTo(unit));
+        }
+    }
 }