fix another test

jimczi · Dec 6, 2024 · 13499d3 · 13499d3
1 parent 28a06f6
commit 13499d3
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/...ce/src/test/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilderTests.java b/...ce/src/test/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilderTests.java
@@ -52,6 +52,7 @@
 import org.elasticsearch.xpack.core.ml.inference.MlInferenceNamedXContentProvider;
 import org.elasticsearch.xpack.core.ml.inference.results.MlTextEmbeddingResults;
 import org.elasticsearch.xpack.core.ml.inference.results.TextExpansionResults;
+import org.elasticsearch.xpack.core.ml.search.SparseVectorQueryWrapper;
 import org.elasticsearch.xpack.core.ml.search.WeightedToken;
 import org.elasticsearch.xpack.inference.InferencePlugin;
 import org.elasticsearch.xpack.inference.mapper.SemanticTextField;
@@ -194,7 +195,7 @@ protected void doAssertLuceneQuery(SemanticQueryBuilder queryBuilder, Query quer
     }
 
     private void assertSparseEmbeddingLuceneQuery(Query query) {
-        Query innerQuery = assertOuterBooleanQuery(query);
+        Query innerQuery = assertOuterSparseVectorQuery(query);
         assertThat(innerQuery, instanceOf(BooleanQuery.class));
 
         BooleanQuery innerBooleanQuery = (BooleanQuery) innerQuery;
@@ -207,7 +208,7 @@ private void assertSparseEmbeddingLuceneQuery(Query query) {
     }
 
     private void assertTextEmbeddingLuceneQuery(Query query) {
-        Query innerQuery = assertOuterBooleanQuery(query);
+        Query innerQuery = assertOuterSparseVectorQuery(query);
 
         Class<? extends Query> expectedKnnQueryClass = switch (denseVectorElementType) {
             case FLOAT -> KnnFloatVectorQuery.class;
@@ -217,9 +218,11 @@ private void assertTextEmbeddingLuceneQuery(Query query) {
         assertThat(innerQuery, instanceOf(expectedKnnQueryClass));
     }
 
-    private Query assertOuterBooleanQuery(Query query) {
-        assertThat(query, instanceOf(BooleanQuery.class));
-        BooleanQuery outerBooleanQuery = (BooleanQuery) query;
+    private Query assertOuterSparseVectorQuery(Query query) {
+        assertThat(query, instanceOf(SparseVectorQueryWrapper.class));
+        var wrapper = (SparseVectorQueryWrapper) query;
+        assertThat(wrapper.getTermsQuery(), instanceOf(BooleanQuery.class));
+        BooleanQuery outerBooleanQuery = (BooleanQuery) wrapper.getTermsQuery();
 
         List<BooleanClause> outerMustClauses = new ArrayList<>();
         List<BooleanClause> outerFilterClauses = new ArrayList<>();