Added unit tests

Signed-off-by: Martin Gaievski <[email protected]>
martin-gaievski · Nov 16, 2024 · 7f20f72 · 7f20f72
1 parent 484d820
commit 7f20f72
Showing 1 changed file with 30 additions and 0 deletions.
diff --git a/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java b/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java
@@ -1197,6 +1197,36 @@ private void createTimeoutCheckerThenWaitThenRun(
         verifyNoMoreInteractions(mockedSearchContext);
     }
 
+    public void testQueryCollectorContextScoreModes() throws Exception {
+        // Test default empty context
+        QueryCollectorContext emptyContext = QueryCollectorContext.EMPTY_CONTEXT;
+        assertEquals(org.apache.lucene.search.ScoreMode.COMPLETE_NO_SCORES, emptyContext.create(null).scoreMode());
+
+        // Test getting context for different score modes
+        org.apache.lucene.search.ScoreMode[] scoreModes = org.apache.lucene.search.ScoreMode.values();
+        for (org.apache.lucene.search.ScoreMode scoreMode : scoreModes) {
+            QueryCollectorContext context = QueryCollectorContext.getContextForScoreMode(scoreMode);
+            assertEquals(scoreMode, context.create(null).scoreMode());
+        }
+
+        // Test that invalid score mode returns empty context
+        QueryCollectorContext defaultContext = QueryCollectorContext.getContextForScoreMode(null);
+        assertEquals(QueryCollectorContext.EMPTY_CONTEXT, defaultContext);
+        assertEquals(org.apache.lucene.search.ScoreMode.COMPLETE_NO_SCORES, defaultContext.create(null).scoreMode());
+    }
+
+    public void testQueryCollectorContextTopScores() throws Exception {
+        QueryCollectorContext topScoresContext = QueryCollectorContext.getContextForScoreMode(
+            org.apache.lucene.search.ScoreMode.TOP_SCORES
+        );
+
+        // Verify score mode
+        assertEquals(org.apache.lucene.search.ScoreMode.TOP_SCORES, topScoresContext.create(null).scoreMode());
+
+        // Verify it's a different instance than empty context
+        assertNotEquals(QueryCollectorContext.EMPTY_CONTEXT, topScoresContext);
+    }
+
     private static class TestSearchContextWithRewriteAndCancellation extends TestSearchContext {
 
         private TestSearchContextWithRewriteAndCancellation(