Handle null totHits

Signed-off-by: Finn Carroll <[email protected]>
finnegancarroll · Aug 28, 2024 · a1a6b94 · a1a6b94
1 parent 7f0b06d
commit a1a6b94
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 21 deletions.
diff --git a/server/src/main/java/org/opensearch/transport/protobuf/SearchHitsProtobuf.java b/server/src/main/java/org/opensearch/transport/protobuf/SearchHitsProtobuf.java
@@ -61,9 +61,10 @@ SearchHitsProto toProto() {
             builder.addHits(new SearchHitProtobuf(hit).toProto());
         }
 
-        TotalHits totHits = totalHits;
-        TotalHitsProto.Builder totHitsBuilder = TotalHitsProto.newBuilder().setRelation(totHits.relation.ordinal()).setValue(totHits.value);
-        builder.setTotalHits(totHitsBuilder);
+        if (totalHits != null) {
+            TotalHitsProto.Builder totHitsBuilder = TotalHitsProto.newBuilder().setRelation(totalHits.relation.ordinal()).setValue(totalHits.value);
+            builder.setTotalHits(totHitsBuilder);
+        }
 
         try (BytesStreamOutput sortOut = new BytesStreamOutput()) {
             sortOut.writeOptionalArray(Lucene::writeSortField, sortFields);

diff --git a/server/src/test/java/org/opensearch/search/SearchHitsProtobufTests.java b/server/src/test/java/org/opensearch/search/SearchHitsProtobufTests.java
@@ -75,24 +75,6 @@ public static SearchHitsProtobuf createTestItem(boolean withOptionalInnerHits, b
         return createTestItem(randomFrom(XContentType.values()), withOptionalInnerHits, withShardTarget);
     }
 
-    private static SearchHit[] createSearchHitArray(
-        int size,
-        final MediaType mediaType,
-        boolean withOptionalInnerHits,
-        boolean transportSerialization
-    ) {
-        SearchHit[] hits = new SearchHit[size];
-        for (int i = 0; i < hits.length; i++) {
-            hits[i] = SearchHitTests.createTestItem(mediaType, withOptionalInnerHits, transportSerialization);
-        }
-        return hits;
-    }
-
-    private static TotalHits randomTotalHits(TotalHits.Relation relation) {
-        long totalHits = TestUtil.nextLong(random(), 0, Long.MAX_VALUE);
-        return new TotalHits(totalHits, relation);
-    }
-
     public static SearchHitsProtobuf createTestItem(final MediaType mediaType, boolean withOptionalInnerHits, boolean transportSerialization) {
         return createTestItem(mediaType, withOptionalInnerHits, transportSerialization, randomFrom(TotalHits.Relation.values()));
     }
@@ -118,6 +100,24 @@ private static SearchHitsProtobuf createTestItem(
         return new SearchHitsProtobuf(new SearchHits(hits, totalHits, maxScore, sortFields, collapseField, collapseValues));
     }
 
+    private static SearchHit[] createSearchHitArray(
+        int size,
+        final MediaType mediaType,
+        boolean withOptionalInnerHits,
+        boolean transportSerialization
+    ) {
+        SearchHit[] hits = new SearchHit[size];
+        for (int i = 0; i < hits.length; i++) {
+            hits[i] = SearchHitTests.createTestItem(mediaType, withOptionalInnerHits, transportSerialization);
+        }
+        return hits;
+    }
+
+    private static TotalHits randomTotalHits(TotalHits.Relation relation) {
+        long totalHits = TestUtil.nextLong(random(), 0, Long.MAX_VALUE);
+        return new TotalHits(totalHits, relation);
+    }
+
     private static SortField[] createSortFields(int size) {
         SortField[] sortFields = new SortField[size];
         for (int i = 0; i < sortFields.length; i++) {