diff --git a/src/test/java/org/elasticsearch/search/highlight/CustomHighlighter.java b/src/test/java/org/elasticsearch/search/highlight/CustomHighlighter.java
index 4f21601c669..2845af198a2 100644
--- a/src/test/java/org/elasticsearch/search/highlight/CustomHighlighter.java
+++ b/src/test/java/org/elasticsearch/search/highlight/CustomHighlighter.java
@@ -50,7 +50,7 @@ public class CustomHighlighter implements Highlighter {
             if (cacheEntry.docId == docId) {
                 cacheEntry.position++;
             } else {
-                cacheEntry.docId = highlighterContext.hitContext.docId();
+                cacheEntry.docId = docId;
                 cacheEntry.position = 1;
             }
         }