add unit tests for errors with retrievers

jdconrad · Feb 2, 2024 · 99ac159 · 99ac159
1 parent 633e762
commit 99ac159
Show file tree

Hide file tree

Showing 4 changed files with 192 additions and 2 deletions.
diff --git a/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java
@@ -52,7 +52,7 @@ public static RetrieverBuilder<?> parseTopLevelRetrieverBuilder(XContentParser p
 
             @Override
             public <T> T namedObject(Class<T> categoryClass, String name, Object context) throws IOException {
-                if (categoryClass.equals(QueryBuilder.class)) {
+                if (categoryClass.equals(RetrieverBuilder.class)) {
                     nestedDepth++;
 
                     if (nestedDepth > 2) {

diff --git a/server/src/test/java/org/elasticsearch/search/retriever/RetrieverBuilderErrorTests.java b/server/src/test/java/org/elasticsearch/search/retriever/RetrieverBuilderErrorTests.java
@@ -0,0 +1,91 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.search.retriever;
+
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.search.SearchModule;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xcontent.NamedXContentRegistry;
+import org.elasticsearch.xcontent.XContentParser;
+import org.elasticsearch.xcontent.json.JsonXContent;
+
+import java.io.IOException;
+import java.util.List;
+
+/**
+ * Tests exceptions related to usage of restricted global values with a retriever.
+ */
+public class RetrieverBuilderErrorTests extends ESTestCase {
+
+    public void testRetrieverExtractionErrors() throws IOException {
+        try (
+            XContentParser parser = createParser(
+                JsonXContent.jsonXContent,
+                "{\"query\": {\"match_all\": {}}, \"retriever\":{\"standard\":{}}}"
+            )
+        ) {
+            SearchSourceBuilder ssb = new SearchSourceBuilder();
+            IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> ssb.parseXContent(parser, true, nf -> true));
+            assertEquals("cannot specify both [query] and [retriever]", iae.getMessage());
+        }
+
+        try (
+            XContentParser parser = createParser(
+                JsonXContent.jsonXContent,
+                "{\"knn\":{\"field\": \"test\", \"k\": 2, \"num_candidates\": 5,"
+                    + " \"query_vector\": [1, 2, 3]}, \"retriever\":{\"standard\":{}}}"
+            )
+        ) {
+            SearchSourceBuilder ssb = new SearchSourceBuilder();
+            IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> ssb.parseXContent(parser, true, nf -> true));
+            assertEquals("cannot specify both [knn] and [retriever]", iae.getMessage());
+        }
+
+        try (XContentParser parser = createParser(JsonXContent.jsonXContent, "{\"search_after\": [1], \"retriever\":{\"standard\":{}}}")) {
+            SearchSourceBuilder ssb = new SearchSourceBuilder();
+            IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> ssb.parseXContent(parser, true, nf -> true));
+            assertEquals("cannot specify both [search_after] and [retriever]", iae.getMessage());
+        }
+
+        try (XContentParser parser = createParser(JsonXContent.jsonXContent, "{\"terminate_after\": 1, \"retriever\":{\"standard\":{}}}")) {
+            SearchSourceBuilder ssb = new SearchSourceBuilder();
+            IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> ssb.parseXContent(parser, true, nf -> true));
+            assertEquals("cannot specify both [terminate_after] and [retriever]", iae.getMessage());
+        }
+
+        try (XContentParser parser = createParser(JsonXContent.jsonXContent, "{\"sort\": [\"field\"], \"retriever\":{\"standard\":{}}}")) {
+            SearchSourceBuilder ssb = new SearchSourceBuilder();
+            IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> ssb.parseXContent(parser, true, nf -> true));
+            assertEquals("cannot specify both [sort] and [retriever]", iae.getMessage());
+        }
+
+        try (
+            XContentParser parser = createParser(
+                JsonXContent.jsonXContent,
+                "{\"rescore\": {\"query\": {\"rescore_query\": {\"match_all\": {}}}}, \"retriever\":{\"standard\":{}}}"
+            )
+        ) {
+            SearchSourceBuilder ssb = new SearchSourceBuilder();
+            IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> ssb.parseXContent(parser, true, nf -> true));
+            assertEquals("cannot specify both [rescore] and [retriever]", iae.getMessage());
+        }
+
+        try (XContentParser parser = createParser(JsonXContent.jsonXContent, "{\"min_score\": 2, \"retriever\":{\"standard\":{}}}")) {
+            SearchSourceBuilder ssb = new SearchSourceBuilder();
+            IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> ssb.parseXContent(parser, true, nf -> true));
+            assertEquals("cannot specify both [min_score] and [retriever]", iae.getMessage());
+        }
+    }
+
+    @Override
+    protected NamedXContentRegistry xContentRegistry() {
+        return new NamedXContentRegistry(new SearchModule(Settings.EMPTY, List.of()).getNamedXContents());
+    }
+}
diff --git a/...k/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilder.java b/...k/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilder.java
@@ -76,7 +76,7 @@ public void doExtractToSearchSourceBuilder(SearchSourceBuilder searchSourceBuild
         if (searchSourceBuilder.rankBuilder() == null) {
             searchSourceBuilder.rankBuilder(new RRFRankBuilder(windowSize, rankConstant));
         } else {
-            throw new IllegalStateException("[rank] cannot be declared on multiple retrievers");
+            throw new IllegalArgumentException("[rank] cannot be declared on multiple retrievers");
         }
     }
 }
diff --git a/...gin/rank-rrf/src/test/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilderTests.java b/...gin/rank-rrf/src/test/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilderTests.java
@@ -23,6 +23,7 @@
 import java.io.IOException;
 import java.util.List;
 
+/** Tests for the rrf retriever. */
 public class RRFRetrieverBuilderTests extends ESTestCase {
 
     /** Tests the rrf retriever validates on its own {@link NodeFeature} */
@@ -37,6 +38,96 @@ public void testRetrieverVersions() throws IOException {
         }
     }
 
+    /** Tests extraction errors related to compound retrievers. These tests require a compound retriever which is why they are here. */
+    public void testRetrieverExtractionErrors() throws IOException {
+        try (
+            XContentParser parser = createParser(
+                JsonXContent.jsonXContent,
+                "{\"retriever\":{\"rrf_nl\":{\"retrievers\":"
+                    + "[{\"standard\":{\"search_after\":[1]}},{\"standard\":{\"search_after\":[2]}}]}}}"
+            )
+        ) {
+            SearchSourceBuilder ssb = new SearchSourceBuilder();
+            IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> ssb.parseXContent(parser, true, nf -> true));
+            assertEquals("[search_after] cannot be declared on multiple retrievers", iae.getMessage());
+        }
+
+        try (
+            XContentParser parser = createParser(
+                JsonXContent.jsonXContent,
+                "{\"retriever\":{\"rrf_nl\":{\"retrievers\":"
+                    + "[{\"standard\":{\"terminate_after\":1}},{\"standard\":{\"terminate_after\":2}}]}}}"
+            )
+        ) {
+            SearchSourceBuilder ssb = new SearchSourceBuilder();
+            IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> ssb.parseXContent(parser, true, nf -> true));
+            assertEquals("[terminate_after] cannot be declared on multiple retrievers", iae.getMessage());
+        }
+
+        try (
+            XContentParser parser = createParser(
+                JsonXContent.jsonXContent,
+                "{\"retriever\":{\"rrf_nl\":{\"retrievers\":" + "[{\"standard\":{\"sort\":[\"f1\"]}},{\"standard\":{\"sort\":[\"f2\"]}}]}}}"
+            )
+        ) {
+            SearchSourceBuilder ssb = new SearchSourceBuilder();
+            IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> ssb.parseXContent(parser, true, nf -> true));
+            assertEquals("[sort] cannot be declared on multiple retrievers", iae.getMessage());
+        }
+
+        try (
+            XContentParser parser = createParser(
+                JsonXContent.jsonXContent,
+                "{\"retriever\":{\"rrf_nl\":{\"retrievers\":" + "[{\"standard\":{\"min_score\":1}},{\"standard\":{\"min_score\":2}}]}}}"
+            )
+        ) {
+            SearchSourceBuilder ssb = new SearchSourceBuilder();
+            IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> ssb.parseXContent(parser, true, nf -> true));
+            assertEquals("[min_score] cannot be declared on multiple retrievers", iae.getMessage());
+        }
+
+        try (
+            XContentParser parser = createParser(
+                JsonXContent.jsonXContent,
+                "{\"retriever\":{\"rrf_nl\":{\"retrievers\":"
+                    + "[{\"standard\":{\"collapse\":{\"field\":\"f0\"}}},{\"standard\":{\"collapse\":{\"field\":\"f1\"}}}]}}}"
+            )
+        ) {
+            SearchSourceBuilder ssb = new SearchSourceBuilder();
+            IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> ssb.parseXContent(parser, true, nf -> true));
+            assertEquals("[collapse] cannot be declared on multiple retrievers", iae.getMessage());
+        }
+
+        try (
+            XContentParser parser = createParser(
+                JsonXContent.jsonXContent,
+                "{\"retriever\":{\"rrf_nl\":{\"retrievers\":[{\"rrf_nl\":{}}]}}}"
+            )
+        ) {
+            SearchSourceBuilder ssb = new SearchSourceBuilder();
+            IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> ssb.parseXContent(parser, true, nf -> true));
+            assertEquals("[rank] cannot be declared on multiple retrievers", iae.getMessage());
+        }
+    }
+
+    /** Tests max depth errors related to compound retrievers. These tests require a compound retriever which is why they are here. */
+    public void testRetrieverBuilderParsingMaxDepth() throws IOException {
+        try (
+            XContentParser parser = createParser(
+                JsonXContent.jsonXContent,
+                "{\"retriever\":{\"rrf_nl\":{\"retrievers\":[{\"rrf_nl\":{\"retrievers\":[{\"standard\":{}}]}}]}}}"
+            )
+        ) {
+            SearchSourceBuilder ssb = new SearchSourceBuilder();
+            IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> ssb.parseXContent(parser, true, nf -> true));
+            assertEquals("[1:65] [rrf] failed to parse field [retrievers]", iae.getMessage());
+            assertEquals(
+                "the nested depth of the [standard] retriever exceeds the maximum nested depth [2] for retrievers",
+                iae.getCause().getCause().getMessage()
+            );
+        }
+    }
+
     @Override
     protected NamedXContentRegistry xContentRegistry() {
         List<NamedXContentRegistry.Entry> entries = new SearchModule(Settings.EMPTY, List.of()).getNamedXContents();
@@ -47,6 +138,14 @@ protected NamedXContentRegistry xContentRegistry() {
                 (p, c) -> RRFRetrieverBuilder.fromXContent(p, (RetrieverParserContext) c)
             )
         );
+        // Add an entry with no license requirement for unit testing
+        entries.add(
+            new NamedXContentRegistry.Entry(
+                RetrieverBuilder.class,
+                new ParseField(RRFRankPlugin.NAME + "_nl"),
+                (p, c) -> RRFRetrieverBuilder.PARSER.apply(p, (RetrieverParserContext) c)
+            )
+        );
         return new NamedXContentRegistry(entries);
     }
 }