Fixed unit test for neural query after recent knn change in rescore c…

…ontext Signed-off-by: Martin Gaievski <[email protected]>
opensearch-project · Oct 9, 2024 · 98fa414 · 98fa414
1 parent 9f4a49a
commit 98fa414
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java
@@ -16,6 +16,7 @@
 import static org.opensearch.knn.index.query.KNNQueryBuilder.MIN_SCORE_FIELD;
 import static org.opensearch.knn.index.query.KNNQueryBuilder.RESCORE_FIELD;
 import static org.opensearch.knn.index.query.KNNQueryBuilder.RESCORE_OVERSAMPLE_FIELD;
+import static org.opensearch.neuralsearch.util.TestUtils.DELTA_FOR_FLOATS_ASSERTION;
 import static org.opensearch.neuralsearch.util.TestUtils.xContentBuilderToMap;
 import static org.opensearch.neuralsearch.query.NeuralQueryBuilder.K_FIELD;
 import static org.opensearch.neuralsearch.query.NeuralQueryBuilder.MODEL_ID_FIELD;
@@ -183,7 +184,11 @@ public void testFromXContent_withRescoreContext_thenBuildSuccessfully() {
         assertEquals(QUERY_TEXT, neuralQueryBuilder.queryText());
         assertEquals(MODEL_ID, neuralQueryBuilder.modelId());
         assertEquals(K, neuralQueryBuilder.k());
-        assertEquals(RescoreContext.getDefault(), neuralQueryBuilder.rescoreContext());
+        assertEquals(
+            RescoreContext.getDefault().getOversampleFactor(),
+            neuralQueryBuilder.rescoreContext().getOversampleFactor(),
+            DELTA_FOR_FLOATS_ASSERTION
+        );
         assertNull(neuralQueryBuilder.methodParameters());
     }
 

diff --git a/src/testFixtures/java/org/opensearch/neuralsearch/util/TestUtils.java b/src/testFixtures/java/org/opensearch/neuralsearch/util/TestUtils.java
@@ -39,6 +39,7 @@ public class TestUtils {
 
     public static final String RELATION_EQUAL_TO = "eq";
     public static final float DELTA_FOR_SCORE_ASSERTION = 0.001f;
+    public static final float DELTA_FOR_FLOATS_ASSERTION = 0.001f;
     public static final String RESTART_UPGRADE_OLD_CLUSTER = "tests.is_old_cluster";
     public static final String BWC_VERSION = "tests.plugin_bwc_version";
     public static final String NEURAL_SEARCH_BWC_PREFIX = "neuralsearch-bwc-";