Fix test

kderusso · Mar 4, 2024 · 0a0ef48 · 0a0ef48
1 parent 2710d23
commit 0a0ef48
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/...n/ml/src/test/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilderTests.java b/...n/ml/src/test/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilderTests.java
@@ -25,6 +25,7 @@
 import org.elasticsearch.common.compress.CompressedXContent;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.extras.MapperExtrasPlugin;
+import org.elasticsearch.index.query.BoolQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.plugins.Plugin;
@@ -259,6 +260,10 @@ public void testThatTokensAreCorrectlyPruned() {
         SearchExecutionContext searchExecutionContext = createSearchExecutionContext();
         TextExpansionQueryBuilder queryBuilder = createTestQueryBuilder();
         QueryBuilder rewrittenQueryBuilder = rewriteAndFetch(queryBuilder, searchExecutionContext);
-        assertTrue(rewrittenQueryBuilder instanceof WeightedTokensQueryBuilder);
+        if (queryBuilder.getTokenPruningConfig() == null) {
+            assertTrue(rewrittenQueryBuilder instanceof BoolQueryBuilder);
+        } else {
+            assertTrue(rewrittenQueryBuilder instanceof WeightedTokensQueryBuilder);
+        }
     }
 }