modify test to ensure that query is rewritten to a supported type

Author: henriquepaes1

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilderTests.java

@@ -283,14 +283,11 @@ protected String[] shuffleProtectedFields() {
         return new String[] { WeightedTokensQueryBuilder.TOKENS_FIELD.getPreferredName() };
     }
 
-    public void testThatTokensAreCorrectlyPruned() {
+    public void testQueryWasRewrittenToASupportedType() {
         SearchExecutionContext searchExecutionContext = createSearchExecutionContext();
         TextExpansionQueryBuilder queryBuilder = createTestQueryBuilder();
         QueryBuilder rewrittenQueryBuilder = rewriteAndFetch(queryBuilder, searchExecutionContext);
-        if (queryBuilder.getTokenPruningConfig() == null) {
-            assertTrue(rewrittenQueryBuilder instanceof BoolQueryBuilder);
-        } else {
-            assertTrue(rewrittenQueryBuilder instanceof WeightedTokensQueryBuilder);
-        }
+
+        assertTrue(rewrittenQueryBuilder instanceof BoolQueryBuilder || rewrittenQueryBuilder instanceof WeightedTokensQueryBuilder);
     }
 }