Fix tests

Author: nik9000

server/src/test/java/org/elasticsearch/index/mapper/DateFieldMapperTests.java

@@ -36,6 +36,7 @@
 import org.elasticsearch.index.mapper.MapperService.MergeReason;
 import org.elasticsearch.index.termvectors.TermVectorsService;
 import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.test.ESSingleNodeTestCase;
 import org.elasticsearch.test.InternalSettingsPlugin;
 import org.junit.Before;
@@ -494,17 +495,19 @@ public void testFetchValuesNanos() throws IOException {
     }
 
     public void testFetchDocValuesMillis() throws IOException {
-        DateFieldMapper mapper = createMapper(Resolution.MILLISECONDS, "strict_date_time");
+        DateFieldMapper mapper = createMapper(Resolution.MILLISECONDS, "strict_date_time||epoch_millis");
+        DocValueFormat format = mapper.mappedFieldType.docValueFormat(null, null);
         String date = "2020-05-15T21:33:02.123Z";
-        assertEquals(List.of(date), fetchFromDocValues(mapper, null, date));
-        assertEquals(List.of(date), fetchFromDocValues(mapper, null, 1589578382123L));
+        assertEquals(List.of(date), fetchFromDocValues(mapper, format, date));
+        assertEquals(List.of(date), fetchFromDocValues(mapper, format, 1589578382123L));
     }
 
     public void testFetchDocValuesNanos() throws IOException {
-        DateFieldMapper mapper = createMapper(Resolution.NANOSECONDS, "strict_date_time");
+        DateFieldMapper mapper = createMapper(Resolution.NANOSECONDS, "strict_date_time||epoch_millis");
+        DocValueFormat format = DocValueFormat.withNanosecondResolution(mapper.mappedFieldType.docValueFormat(null, null));
         String date = "2020-05-15T21:33:02.123456789Z";
-        assertEquals(List.of(date), fetchFromDocValues(mapper, null, date));
-        assertEquals(List.of("2020-05-15T21:33:02.123Z"), fetchFromDocValues(mapper, null, 1589578382123L));
+        assertEquals(List.of(date), fetchFromDocValues(mapper, format, date));
+        assertEquals(List.of("2020-05-15T21:33:02.123Z"), fetchFromDocValues(mapper, format, 1589578382123L));
     }
 
     private DateFieldMapper createMapper(Resolution resolution, String format) {

test/framework/src/main/java/org/elasticsearch/test/ESSingleNodeTestCase.java

@@ -30,7 +30,6 @@
 import org.apache.lucene.search.Scorable;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.store.Directory;
-import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.SetOnce;
 import org.elasticsearch.action.admin.cluster.health.ClusterHealthResponse;
 import org.elasticsearch.action.admin.indices.create.CreateIndexRequestBuilder;
@@ -45,8 +44,7 @@
 import org.elasticsearch.cluster.routing.allocation.DiskThresholdSettings;
 import org.elasticsearch.common.Priority;
 import org.elasticsearch.common.Strings;
-import org.elasticsearch.common.bytes.BytesArray;
-import org.elasticsearch.common.lucene.BytesRefs;
+import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.common.util.BigArrays;
@@ -55,6 +53,8 @@
 import org.elasticsearch.common.xcontent.NamedXContentRegistry;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.common.xcontent.XContentParser.Token;
+import org.elasticsearch.common.xcontent.XContentParserUtils;
 import org.elasticsearch.common.xcontent.XContentType;
 import org.elasticsearch.common.xcontent.json.JsonXContent;
 import org.elasticsearch.core.internal.io.IOUtils;
@@ -444,36 +444,40 @@ public void collect(int doc) throws IOException {
     /**
      * Use a {@linkplain FieldMapper} to extract values from doc values.
      */
-    protected static List<?> fetchFromDocValues(FieldMapper mapper, String format, Object sourceValue) throws IOException {
+    protected static List<?> fetchFromDocValues(FieldMapper mapper, DocValueFormat format, Object sourceValue) throws IOException {
         MapperService mapperService = mock(MapperService.class);
         when(mapperService.fieldType(any())).thenReturn(mapper.fieldType());
         Function<MappedFieldType, IndexFieldData<?>> fieldDataLookup = mft -> mft.fielddataBuilder("test", () -> {
             throw new UnsupportedOperationException();
         }).build(new IndexFieldDataCache.None(), new NoneCircuitBreakerService(), mapperService);
         try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
-            BytesRef source = BytesRefs.toBytesRef(
+            BytesReference source = BytesReference.bytes(
                 JsonXContent.contentBuilder().startObject().field(mapper.name(), sourceValue).endObject()
             );
-            ParseContext.Document doc = new ParseContext.Document();
-            ParseContext context = mock(ParseContext.class);
-            when(context.doc()).thenReturn(doc);
-            when(context.sourceToParse()).thenReturn(new SourceToParse("test", "id", new BytesArray(source), XContentType.JSON));
             try (
                 XContentParser parser = JsonXContent.jsonXContent.createParser(
                     NamedXContentRegistry.EMPTY,
                     DeprecationHandler.THROW_UNSUPPORTED_OPERATION,
-                    source.utf8ToString()
+                    source.streamInput()
                 )
             ) {
+                ParseContext.Document doc = new ParseContext.Document();
+                ParseContext context = mock(ParseContext.class);
+                when(context.doc()).thenReturn(doc);
+                when(context.sourceToParse()).thenReturn(new SourceToParse("test", "id", source, XContentType.JSON));
+                when(context.parser()).thenReturn(parser);
                 parser.nextToken();
+                XContentParserUtils.ensureExpectedToken(Token.START_OBJECT, parser.currentToken(), parser::getTokenLocation);
+                XContentParserUtils.ensureFieldName(parser, parser.nextToken(), mapper.name());
                 parser.nextToken();
                 mapper.parse(context);
+                parser.nextToken();
+                XContentParserUtils.ensureExpectedToken(Token.END_OBJECT, parser.currentToken(), parser::getTokenLocation);
+                iw.addDocument(doc);
             }
-            iw.addDocument(doc);
             try (DirectoryReader reader = iw.getReader()) {
                 IndexSearcher indexSearcher = newSearcher(reader);
                 SearchLookup lookup = new SearchLookup(mapperService, fieldDataLookup);
-                DocValueFormat dvFormat = mapper.fieldType().docValueFormat(format, null);
                 IndexFieldData<?> ifd = lookup.doc().getForField(mapper.fieldType());
                 SetOnce<List<?>> result = new SetOnce<>();
                 indexSearcher.search(new MatchAllDocsQuery(), new Collector() {
@@ -484,7 +488,7 @@ public ScoreMode scoreMode() {
 
                     @Override
                     public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
-                        LeafValueFetcher lvf = ifd.load(context).buildFetcher(dvFormat);
+                        LeafValueFetcher lvf = ifd.load(context).buildFetcher(format);
                         return new LeafCollector() {
                             @Override
                             public void setScorer(Scorable scorer) throws IOException {}