Fixed rewrite tests

Author: Mikep86

server/src/main/java/org/elasticsearch/search/retriever/CompoundRetrieverBuilder.java

@@ -36,6 +36,7 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
 import java.util.Objects;
@@ -295,6 +296,10 @@ public int rankWindowSize() {
         return rankWindowSize;
     }
 
+    public List<RetrieverSource> innerRetrievers() {
+        return Collections.unmodifiableList(innerRetrievers);
+    }
+
     public static RetrieverSource convertToRetrieverSource(RetrieverBuilder retrieverBuilder) {
         return new RetrieverSource(retrieverBuilder, null);
     }

x-pack/plugin/rank-rrf/src/test/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilderTests.java

@@ -34,8 +34,10 @@
 import org.elasticsearch.xcontent.json.JsonXContent;
 
 import java.io.IOException;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import static org.elasticsearch.search.retriever.CompoundRetrieverBuilder.convertToRetrieverSource;
 
@@ -164,28 +166,36 @@ private static void assertSimplifiedParamsRewrite(
         Map<String, Float> expectedInferenceFields,
         String expectedQuery
     ) {
-        List<CompoundRetrieverBuilder.RetrieverSource> expectedInnerRetrievers = List.of(
+        Set<Object> expectedInnerRetrievers = Set.of(
             convertToRetrieverSource(
                 new StandardRetrieverBuilder(
                     new MultiMatchQueryBuilder(expectedQuery).type(MultiMatchQueryBuilder.Type.MOST_FIELDS)
                         .fields(expectedNonInferenceFields)
                 )
             ),
-            convertToRetrieverSource(new RRFRetrieverBuilder(expectedInferenceFields.entrySet().stream().map(e -> {
+            Set.of(expectedInferenceFields.entrySet().stream().map(e -> {
                 if (e.getValue() != 1.0f) {
                     throw new IllegalArgumentException("Cannot apply per-field weights in RRF");
                 }
                 return convertToRetrieverSource(new StandardRetrieverBuilder(new MatchQueryBuilder(e.getKey(), expectedQuery)));
-            }).toList(), retriever.rankWindowSize(), retriever.rankConstant()))
-        );
-        RRFRetrieverBuilder expectedRewritten = new RRFRetrieverBuilder(
-            expectedInnerRetrievers,
-            retriever.rankWindowSize(),
-            retriever.rankConstant()
+            }).toArray())
         );
 
         RRFRetrieverBuilder rewritten = retriever.doRewrite(ctx);
         assertNotSame(retriever, rewritten);
-        assertEquals(expectedRewritten, rewritten);
+        assertEquals(expectedInnerRetrievers, getInnerRetrieversAsSet(rewritten));
+    }
+
+    private static Set<Object> getInnerRetrieversAsSet(RRFRetrieverBuilder retriever) {
+        Set<Object> innerRetrieversSet = new HashSet<>();
+        for (CompoundRetrieverBuilder.RetrieverSource innerRetriever : retriever.innerRetrievers()) {
+            if (innerRetriever.retriever() instanceof RRFRetrieverBuilder innerRrfRetriever) {
+                innerRetrieversSet.add(getInnerRetrieversAsSet(innerRrfRetriever));
+            } else {
+                innerRetrieversSet.add(innerRetriever);
+            }
+        }
+
+        return innerRetrieversSet;
     }
 }