Fixed flaky unit test

Signed-off-by: Martin Gaievski <[email protected]>
opensearch-project · Nov 15, 2023 · a2c0593 · a2c0593
1 parent dff866d
commit a2c0593
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 2 deletions.
diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryScorerTests.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryScorerTests.java
@@ -169,6 +169,40 @@ public void testWithRandomDocuments_whenMultipleScorersAndSomeScorersEmpty_thenR
         testWithQuery(docs, scores, hybridQueryScorer);
     }
 
+    @SneakyThrows
+    public void testMaxScore_whenMultipleScorers_thenSuccessful() {
+        int maxDocId = TestUtil.nextInt(random(), 10, 10_000);
+        Pair<int[], float[]> docsAndScores = generateDocuments(maxDocId);
+        int[] docs = docsAndScores.getLeft();
+        float[] scores = docsAndScores.getRight();
+
+        Weight weight = mock(Weight.class);
+
+        HybridQueryScorer hybridQueryScorerWithAllNonNullSubScorers = new HybridQueryScorer(
+            weight,
+            Arrays.asList(
+                scorer(docs, scores, fakeWeight(new MatchAllDocsQuery())),
+                scorer(docs, scores, fakeWeight(new MatchNoDocsQuery()))
+            )
+        );
+
+        float maxScore = hybridQueryScorerWithAllNonNullSubScorers.getMaxScore(Integer.MAX_VALUE);
+        assertTrue(maxScore > 0.0f);
+
+        HybridQueryScorer hybridQueryScorerWithSomeNullSubScorers = new HybridQueryScorer(
+            weight,
+            Arrays.asList(null, scorer(docs, scores, fakeWeight(new MatchAllDocsQuery())), null)
+        );
+
+        maxScore = hybridQueryScorerWithSomeNullSubScorers.getMaxScore(Integer.MAX_VALUE);
+        assertTrue(maxScore > 0.0f);
+
+        HybridQueryScorer hybridQueryScorerWithAllNullSubScorers = new HybridQueryScorer(weight, Arrays.asList(null, null));
+
+        maxScore = hybridQueryScorerWithAllNullSubScorers.getMaxScore(Integer.MAX_VALUE);
+        assertEquals(0.0f, maxScore, 0.0f);
+    }
+
     private Pair<int[], float[]> generateDocuments(int maxDocId) {
         final int numDocs = RandomizedTest.randomIntBetween(1, maxDocId / 2);
         final int[] docs = new int[numDocs];

diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryWeightTests.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryWeightTests.java
@@ -61,11 +61,11 @@ public void testScorerIterator_whenExecuteQuery_thenScorerIteratorSuccessful() {
             List.of(QueryBuilders.termQuery(TEXT_FIELD_NAME, TERM_QUERY_TEXT).toQuery(mockQueryShardContext))
         );
         IndexSearcher searcher = newSearcher(reader);
-        Weight weight = searcher.createWeight(hybridQueryWithTerm, ScoreMode.COMPLETE, 1.0f);
+        Weight weight = hybridQueryWithTerm.createWeight(searcher, ScoreMode.TOP_SCORES, 1.0f);
 
         assertNotNull(weight);
 
-        LeafReaderContext leafReaderContext = reader.getContext().leaves().get(0);
+        LeafReaderContext leafReaderContext = searcher.getIndexReader().leaves().get(0);
         Scorer scorer = weight.scorer(leafReaderContext);
 
         assertNotNull(scorer);