improve FetchSourceContext and add unit tests (elastic#77346) (elasti…

…c#77521) Co-authored-by: Elastic Machine <[email protected]> Co-authored-by: weizijun <[email protected]> Co-authored-by: Elastic Machine <[email protected]>
gwbrown · Sep 9, 2021 · c6180cf · c6180cf
1 parent 7da3b15
commit c6180cf
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 7 deletions.
diff --git a/server/src/main/java/org/elasticsearch/search/fetch/subphase/FetchSourceContext.java b/server/src/main/java/org/elasticsearch/search/fetch/subphase/FetchSourceContext.java
@@ -8,17 +8,17 @@
 
 package org.elasticsearch.search.fetch.subphase;
 
-import org.elasticsearch.core.Booleans;
-import org.elasticsearch.common.xcontent.ParseField;
 import org.elasticsearch.common.ParsingException;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.common.xcontent.ParseField;
 import org.elasticsearch.common.xcontent.ToXContentObject;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.common.xcontent.support.XContentMapValues;
+import org.elasticsearch.core.Booleans;
 import org.elasticsearch.rest.RestRequest;
 
 import java.io.IOException;
@@ -111,6 +111,10 @@ public static FetchSourceContext parseFromRestRequest(RestRequest request) {
     }
 
     public static FetchSourceContext fromXContent(XContentParser parser) throws IOException {
+        if (parser.currentToken() == null) {
+            parser.nextToken();
+        }
+
         XContentParser.Token token = parser.currentToken();
         boolean fetchSource = true;
         String[] includes = Strings.EMPTY_ARRAY;
@@ -121,7 +125,7 @@ public static FetchSourceContext fromXContent(XContentParser parser) throws IOEx
             includes = new String[]{parser.text()};
         } else if (token == XContentParser.Token.START_ARRAY) {
             ArrayList<String> list = new ArrayList<>();
-            while ((token = parser.nextToken()) != XContentParser.Token.END_ARRAY) {
+            while ((parser.nextToken()) != XContentParser.Token.END_ARRAY) {
                 list.add(parser.text());
             }
             includes = list.toArray(new String[list.size()]);
@@ -172,8 +176,21 @@ public static FetchSourceContext fromXContent(XContentParser parser) throws IOEx
                 }
             }
         } else {
-            throw new ParsingException(parser.getTokenLocation(), "Expected one of [" + XContentParser.Token.VALUE_BOOLEAN + ", "
-                    + XContentParser.Token.START_OBJECT + "] but found [" + token + "]", parser.getTokenLocation());
+            throw new ParsingException(
+                parser.getTokenLocation(),
+                "Expected one of ["
+                    + XContentParser.Token.VALUE_BOOLEAN
+                    + ", "
+                    + XContentParser.Token.VALUE_STRING
+                    + ", "
+                    + XContentParser.Token.START_ARRAY
+                    + ", "
+                    + XContentParser.Token.START_OBJECT
+                    + "] but found ["
+                    + token
+                    + "]",
+                parser.getTokenLocation()
+            );
         }
         return new FetchSourceContext(fetchSource, includes, excludes);
     }

diff --git a/server/src/test/java/org/elasticsearch/search/fetch/subphase/FetchSourceContextTests.java b/server/src/test/java/org/elasticsearch/search/fetch/subphase/FetchSourceContextTests.java
@@ -0,0 +1,54 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.search.fetch.subphase;
+
+import org.elasticsearch.common.ParsingException;
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.common.xcontent.XContentFactory;
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.test.AbstractSerializingTestCase;
+
+import java.io.IOException;
+
+import static org.hamcrest.Matchers.containsString;
+
+public class FetchSourceContextTests extends AbstractSerializingTestCase<FetchSourceContext> {
+    @Override
+    protected FetchSourceContext doParseInstance(XContentParser parser) throws IOException {
+        return FetchSourceContext.fromXContent(parser);
+    }
+
+    @Override
+    protected Writeable.Reader<FetchSourceContext> instanceReader() {
+        return FetchSourceContext::new;
+    }
+
+    @Override
+    protected FetchSourceContext createTestInstance() {
+        return new FetchSourceContext(
+            true,
+            randomArray(0, 5, String[]::new, () -> randomAlphaOfLength(5)),
+            randomArray(0, 5, String[]::new, () -> randomAlphaOfLength(5))
+        );
+    }
+
+    public void testFromXContentException() throws IOException {
+        XContentBuilder builder = XContentFactory.jsonBuilder();
+        int value = randomInt();
+        builder.value(value);
+        XContentParser parser = createParser(builder);
+        ParsingException exception = expectThrows(ParsingException.class, () -> FetchSourceContext.fromXContent(parser));
+        assertThat(
+            exception.getMessage(),
+            containsString("Expected one of [VALUE_BOOLEAN, VALUE_STRING, START_ARRAY, START_OBJECT] but found [VALUE_NUMBER]")
+        );
+
+    }
+}
diff --git a/test/framework/src/main/java/org/elasticsearch/test/AbstractXContentTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/AbstractXContentTestCase.java
@@ -260,8 +260,8 @@ static BytesReference insertRandomFieldsAndShuffle(BytesReference xContent, XCon
         } else {
             withRandomFields = xContent;
         }
-        XContentParser parserWithRandonFields = createParserFunction.apply(XContentFactory.xContent(xContentType), withRandomFields);
-        return BytesReference.bytes(ESTestCase.shuffleXContent(parserWithRandonFields, false, shuffleFieldsExceptions));
+        XContentParser parserWithRandomFields = createParserFunction.apply(XContentFactory.xContent(xContentType), withRandomFields);
+        return BytesReference.bytes(ESTestCase.shuffleXContent(parserWithRandomFields, false, shuffleFieldsExceptions));
     }
 
 }