From b3ce44cd0d76e1f437490c17a6cf37a7657a0fe9 Mon Sep 17 00:00:00 2001
From: Adrien Grand <jpountz@gmail.com>
Date: Fri, 3 Sep 2021 14:43:32 +0200
Subject: [PATCH] LUCENE-9620: Implement AssertingWeight#count.

---
 .../java/org/apache/lucene/search/AssertingWeight.java   | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/lucene/test-framework/src/java/org/apache/lucene/search/AssertingWeight.java b/lucene/test-framework/src/java/org/apache/lucene/search/AssertingWeight.java
index e61224cf717..857e2dc954e 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/search/AssertingWeight.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/search/AssertingWeight.java
@@ -33,6 +33,15 @@ class AssertingWeight extends FilterWeight {
     this.scoreMode = scoreMode;
   }
 
+  @Override
+  public int count(LeafReaderContext context) throws IOException {
+    final int count = in.count(context);
+    if (count < -1 || count > context.reader().numDocs()) {
+      throw new AssertionError("count=" + count + ", numDocs=" + context.reader().numDocs());
+    }
+    return count;
+  }
+
   @Override
   public Matches matches(LeafReaderContext context, int doc) throws IOException {
     Matches matches = in.matches(context, doc);