add back awaitfix tests

2 files changed, 42 insertions, 44 deletions

diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoContextMapping.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoContextMapping.java
index f895f280d2..149bfa665c 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoContextMapping.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoContextMapping.java
@@ -243,21 +243,30 @@ public class GeoContextMapping extends ContextMapping {
             }
         }
         List<QueryContext> queryContextList = new ArrayList<>();
-        for (GeoQueryContext geoQueryContext : queryContexts) {
-            int minPrecision = Math.min(this.precision, geoQueryContext.precision);
-            int precision = Math.min(minPrecision, geoQueryContext.geoHash.length());
-            String truncatedGeohash = geoQueryContext.geoHash.toString().substring(0, precision);
-            queryContextList.add(new QueryContext(truncatedGeohash, geoQueryContext.boost, truncatedGeohash.length() < this.precision));
-            for (int neighboursPrecision : geoQueryContext.neighbours) {
-                int neighbourPrecision = Math.min(neighboursPrecision, truncatedGeohash.length());
-                String neighbourGeohash = truncatedGeohash.substring(0, neighbourPrecision);
-                Collection<String> locations = new ArrayList<>();
-                GeoHashUtils.addNeighbors(neighbourGeohash, neighbourPrecision, locations);
-                boolean isPrefix = neighbourPrecision < this.precision;
-                for (String location : locations) {
-                    queryContextList.add(new QueryContext(location, geoQueryContext.boost, isPrefix));
+        for (GeoQueryContext queryContext : queryContexts) {
+            int minPrecision = this.precision;
+            if (queryContext.precision != -1) {
+                minPrecision = Math.min(minPrecision, queryContext.precision);
+            }
+            GeoPoint point = queryContext.geoPoint;
+            final Collection<String> locations = new HashSet<>();
+            String geoHash = GeoHashUtils.stringEncode(point.getLon(), point.getLat(), minPrecision);
+            locations.add(geoHash);
+            if (queryContext.neighbours.isEmpty() && geoHash.length() == this.precision) {
+                GeoHashUtils.addNeighbors(geoHash, locations);
+            } else if (queryContext.neighbours.isEmpty() == false) {
+                for (Integer neighbourPrecision : queryContext.neighbours) {
+                    if (neighbourPrecision < geoHash.length()) {
+                        String truncatedGeoHash = geoHash.substring(0, neighbourPrecision);
+                        locations.add(truncatedGeoHash);
+                        GeoHashUtils.addNeighbors(truncatedGeoHash, locations);
+                    }
+
                 }
             }
+            for (String location : locations) {
+                queryContextList.add(new QueryContext(location, queryContext.boost, location.length() < this.precision));
+            }
         }
         return queryContextList;
     }
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoQueryContext.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoQueryContext.java
index 19ea55d804..291f65091e 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoQueryContext.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoQueryContext.java
@@ -29,6 +29,8 @@ import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 import static org.elasticsearch.search.suggest.completion.context.GeoContextMapping.*;
@@ -37,10 +39,10 @@ import static org.elasticsearch.search.suggest.completion.context.GeoContextMapp
  * Defines the query context for {@link GeoContextMapping}
  */
 public final class GeoQueryContext implements ToXContent {
-    public CharSequence geoHash;
+    public GeoPoint geoPoint;
     public int boost = 1;
-    public int precision = DEFAULT_PRECISION;
-    public int[] neighbours;
+    public int precision = -1;
+    public List<Integer> neighbours = new ArrayList<>(0);
 
     /**
      * Creates a query context for a given geo point with a boost of 1
@@ -71,7 +73,7 @@ public final class GeoQueryContext implements ToXContent {
      * provided boost
      */
     public GeoQueryContext(CharSequence geoHash, int boost) {
-        this(geoHash, boost, DEFAULT_PRECISION);
+        this(geoHash, boost, -1);
     }
 
     /**
@@ -79,8 +81,8 @@ public final class GeoQueryContext implements ToXContent {
      * a provided boost and enables generating neighbours
      * at specified precisions
      */
-    public GeoQueryContext(GeoPoint geoPoint, int boost, int precision, int... neighbours) {
-        this(geoPoint.geohash(), boost, precision, neighbours);
+    public GeoQueryContext(CharSequence geoHash, int boost, int precision, Integer... neighbours) {
+        this(GeoPoint.fromGeohash(geoHash.toString()), boost, precision, neighbours);
     }
 
     /**
@@ -88,11 +90,11 @@ public final class GeoQueryContext implements ToXContent {
      * a provided boost and enables generating neighbours
      * at specified precisions
      */
-    public GeoQueryContext(CharSequence geoHash, int boost, int precision, int... neighbours) {
-        this.geoHash = geoHash;
+    public GeoQueryContext(GeoPoint geoPoint, int boost, int precision, Integer... neighbours) {
+        this.geoPoint = geoPoint;
         this.boost = boost;
         this.precision = precision;
-        this.neighbours = neighbours;
+        Collections.addAll(this.neighbours, neighbours);
     }
 
     private GeoQueryContext() {
@@ -107,16 +109,11 @@ public final class GeoQueryContext implements ToXContent {
     }
 
     void setNeighbours(List<Integer> neighbours) {
-        int[] neighbourArray = new int[neighbours.size()];
-        for (int i = 0; i < neighbours.size(); i++) {
-            neighbourArray[i] = neighbours.get(i);
-        }
-        this.neighbours = neighbourArray;
+        this.neighbours = neighbours;
     }
 
-    private GeoPoint point;
-    void setPoint(GeoPoint point) {
-        this.point = point;
+    void setGeoPoint(GeoPoint geoPoint) {
+        this.geoPoint = geoPoint;
     }
 
     private double lat = Double.NaN;
@@ -129,28 +126,19 @@ public final class GeoQueryContext implements ToXContent {
         this.lon = lon;
     }
 
-    void setGeoHash(String geoHash) {
-        this.geoHash = geoHash;
-    }
-
     void finish() {
-        if (point == null) {
+        if (geoPoint == null) {
             if (Double.isNaN(lat) == false && Double.isNaN(lon) == false) {
-                point = new GeoPoint(lat, lon);
+                geoPoint = new GeoPoint(lat, lon);
             } else {
                 throw new ElasticsearchParseException("no geohash or geo point provided");
             }
         }
-        this.geoHash = point.geohash();
-        if (this.neighbours == null) {
-            this.neighbours = new int[]{precision};
-        }
     }
 
     private static ObjectParser<GeoQueryContext, GeoContextMapping> GEO_CONTEXT_PARSER = new ObjectParser<>("geo", null);
     static {
-        GEO_CONTEXT_PARSER.declareField((parser, geoQueryContext, geoContextMapping) -> geoQueryContext.setPoint(GeoUtils.parseGeoPoint(parser)), new ParseField("context"), ObjectParser.ValueType.OBJECT);
-        GEO_CONTEXT_PARSER.declareString(GeoQueryContext::setGeoHash, new ParseField("context"));
+        GEO_CONTEXT_PARSER.declareField((parser, geoQueryContext, geoContextMapping) -> geoQueryContext.setGeoPoint(GeoUtils.parseGeoPoint(parser)), new ParseField("context"), ObjectParser.ValueType.OBJECT);
         GEO_CONTEXT_PARSER.declareInt(GeoQueryContext::setBoost, new ParseField("boost"));
         // TODO : add string support for precision for GeoUtils.geoHashLevelsForPrecision()
         GEO_CONTEXT_PARSER.declareInt(GeoQueryContext::setPrecision, new ParseField("precision"));
@@ -166,7 +154,7 @@ public final class GeoQueryContext implements ToXContent {
         if (token == XContentParser.Token.START_OBJECT) {
             GEO_CONTEXT_PARSER.parse(parser, queryContext);
         } else if (token == XContentParser.Token.VALUE_STRING) {
-            queryContext.setPoint(GeoPoint.fromGeohash(parser.text()));
+            queryContext.setGeoPoint(GeoPoint.fromGeohash(parser.text()));
         } else {
             throw new ElasticsearchParseException("geo context must be an object or string");
         }
@@ -178,7 +166,8 @@ public final class GeoQueryContext implements ToXContent {
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
         builder.startObject(CONTEXT_VALUE);
-        builder.field("geohash", geoHash);
+        builder.field("lat", geoPoint.getLat());
+        builder.field("lon", geoPoint.getLon());
         builder.endObject();
         builder.field(CONTEXT_BOOST, boost);
         builder.field(CONTEXT_NEIGHBOURS, neighbours);