From 7f20f72cdcba5854343f4f361c93b2d2311f1bfc Mon Sep 17 00:00:00 2001
From: Martin Gaievski <gaievski@amazon.com>
Date: Fri, 15 Nov 2024 16:22:04 -0800
Subject: [PATCH] Added unit tests

Signed-off-by: Martin Gaievski <gaievski@amazon.com>
---
 .../search/query/QueryPhaseTests.java         | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java b/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java
index 84057ab1a1b15..f5b5dc346068e 100644
--- 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(