Skip to content

Commit 4c6f284

Browse files
mridula-s109mridula-s109
committed
Added retriever build test file
1 parent 84b6ed7 commit 4c6f284

File tree

3 files changed

+175
-6
lines changed

3 files changed

+175
-6
lines changed

‎x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/SpecifiedDocument.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
9494
return builder.endObject();
9595
}
9696

97-
static final ConstructingObjectParser<SpecifiedDocument, Void> PARSER = new ConstructingObjectParser<>(
97+
public static final ConstructingObjectParser<SpecifiedDocument, Void> PARSER = new ConstructingObjectParser<>(
9898
NAME,
9999
a -> new SpecifiedDocument((String) a[0], (String) a[1])
100100
);

‎x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/retriever/PinnedRetrieverBuilder.java‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -58,11 +58,7 @@ public final class PinnedRetrieverBuilder extends CompoundRetrieverBuilder<Pinne
5858

5959
static {
6060
PARSER.declareStringArray(optionalConstructorArg(), IDS_FIELD);
61-
PARSER.declareObjectArray(optionalConstructorArg(), (p, c) -> {
62-
String index = p.textOrNull();
63-
String id = p.text();
64-
return new SpecifiedDocument(index, id);
65-
}, DOCS_FIELD);
61+
PARSER.declareObjectArray(optionalConstructorArg(), (p, c) -> SpecifiedDocument.PARSER.apply(p, null), DOCS_FIELD);
6662
PARSER.declareNamedObject(constructorArg(), (p, c, n) -> {
6763
RetrieverBuilder innerRetriever = p.namedObject(RetrieverBuilder.class, n, c);
6864
c.trackRetrieverUsage(innerRetriever.getName());

‎x-pack/plugin/search-business-rules/src/test/java/org/elasticsearch/xpack/searchbusinessrules/retriever/PinnedRetrieverBuilderTests.java‎

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.searchbusinessrules.retriever;
9+
10+
import org.elasticsearch.common.Strings;
11+
import org.elasticsearch.common.settings.Settings;
12+
import org.elasticsearch.index.query.QueryBuilder;
13+
import org.elasticsearch.index.query.QueryBuilders;
14+
import org.elasticsearch.search.SearchModule;
15+
import org.elasticsearch.search.builder.SearchSourceBuilder;
16+
import org.elasticsearch.search.retriever.RetrieverBuilder;
17+
import org.elasticsearch.search.retriever.RetrieverParserContext;
18+
import org.elasticsearch.search.retriever.TestRetrieverBuilder;
19+
import org.elasticsearch.test.AbstractXContentTestCase;
20+
import org.elasticsearch.usage.SearchUsage;
21+
import org.elasticsearch.usage.SearchUsageHolder;
22+
import org.elasticsearch.usage.UsageService;
23+
import org.elasticsearch.xcontent.NamedXContentRegistry;
24+
import org.elasticsearch.xcontent.ParseField;
25+
import org.elasticsearch.xcontent.XContentParser;
26+
import org.elasticsearch.xcontent.json.JsonXContent;
27+
import org.elasticsearch.xpack.searchbusinessrules.PinnedQueryBuilder;
28+
import org.elasticsearch.xpack.searchbusinessrules.SpecifiedDocument;
29+
30+
import java.io.IOException;
31+
import java.lang.reflect.Field;
32+
import java.util.ArrayList;
33+
import java.util.List;
34+
35+
import static org.elasticsearch.search.rank.RankBuilder.DEFAULT_RANK_WINDOW_SIZE;
36+
import static org.hamcrest.Matchers.equalTo;
37+
import static org.hamcrest.Matchers.instanceOf;
38+
39+
public class PinnedRetrieverBuilderTests extends AbstractXContentTestCase<PinnedRetrieverBuilder> {
40+
41+
public static PinnedRetrieverBuilder createRandomPinnedRetrieverBuilder() {
42+
boolean useIds = randomBoolean();
43+
boolean useDocs = !useIds || randomBoolean();
44+
45+
List<String> ids = useIds ? List.of(randomAlphaOfLengthBetween(5, 10), randomAlphaOfLengthBetween(5, 10)) : new ArrayList<>();
46+
List<SpecifiedDocument> docs = useDocs ? List.of(
47+
new SpecifiedDocument(randomAlphaOfLengthBetween(5, 10), randomAlphaOfLengthBetween(5, 10)),
48+
new SpecifiedDocument(randomAlphaOfLengthBetween(5, 10), randomAlphaOfLengthBetween(5, 10))
49+
) : new ArrayList<>();
50+
51+
return new PinnedRetrieverBuilder(
52+
ids,
53+
docs,
54+
TestRetrieverBuilder.createRandomTestRetrieverBuilder(),
55+
randomIntBetween(1, 100)
56+
);
57+
}
58+
59+
@Override
60+
protected PinnedRetrieverBuilder createTestInstance() {
61+
return createRandomPinnedRetrieverBuilder();
62+
}
63+
64+
@Override
65+
protected PinnedRetrieverBuilder doParseInstance(XContentParser parser) throws IOException {
66+
return (PinnedRetrieverBuilder) RetrieverBuilder.parseTopLevelRetrieverBuilder(
67+
parser,
68+
new RetrieverParserContext(
69+
new SearchUsage(),
70+
nf -> false
71+
)
72+
);
73+
}
74+
75+
@Override
76+
protected boolean supportsUnknownFields() {
77+
return false;
78+
}
79+
80+
@Override
81+
protected String[] getShuffleFieldsExceptions() {
82+
return new String[] {
83+
PinnedRetrieverBuilder.IDS_FIELD.getPreferredName(),
84+
PinnedRetrieverBuilder.DOCS_FIELD.getPreferredName() };
85+
}
86+
87+
@Override
88+
protected NamedXContentRegistry xContentRegistry() {
89+
List<NamedXContentRegistry.Entry> entries = new SearchModule(Settings.EMPTY, List.of()).getNamedXContents();
90+
entries.add(
91+
new NamedXContentRegistry.Entry(
92+
RetrieverBuilder.class,
93+
TestRetrieverBuilder.TEST_SPEC.getName(),
94+
(p, c) -> TestRetrieverBuilder.TEST_SPEC.getParser().fromXContent(p, (RetrieverParserContext) c),
95+
TestRetrieverBuilder.TEST_SPEC.getName().getForRestApiVersion()
96+
)
97+
);
98+
entries.add(
99+
new NamedXContentRegistry.Entry(
100+
RetrieverBuilder.class,
101+
new ParseField(PinnedRetrieverBuilder.NAME),
102+
(p, c) -> PinnedRetrieverBuilder.PARSER.apply(p, (RetrieverParserContext) c)
103+
)
104+
);
105+
return new NamedXContentRegistry(entries);
106+
}
107+
108+
public void testParserDefaults() throws IOException {
109+
// Inner retriever content only sent to parser
110+
String json = """
111+
{
112+
"ids": [ "id1", "id2" ],
113+
"retriever": { "standard": { "query": { "query_string": { "query": "i like pugs" } } } }
114+
}""";
115+
116+
try (XContentParser parser = createParser(JsonXContent.jsonXContent, json)) {
117+
PinnedRetrieverBuilder parsed = PinnedRetrieverBuilder.PARSER.parse(
118+
parser,
119+
new RetrieverParserContext(new SearchUsage(), nf -> true)
120+
);
121+
assertEquals(DEFAULT_RANK_WINDOW_SIZE, parsed.rankWindowSize());
122+
}
123+
}
124+
125+
public void testPinnedRetrieverParsing() throws IOException {
126+
String restContent = """
127+
{
128+
"retriever": {
129+
"pinned": {
130+
"retriever": {
131+
"test": {
132+
"value": "my-test-retriever"
133+
}
134+
},
135+
"ids": [
136+
"id1",
137+
"id2"
138+
],
139+
"docs": [
140+
{
141+
"_index": "index1",
142+
"_id": "doc1"
143+
},
144+
{
145+
"_index": "index2",
146+
"_id": "doc2"
147+
}
148+
],
149+
"rank_window_size": 100,
150+
"_name": "my_pinned_retriever"
151+
}
152+
}
153+
}""";
154+
SearchUsageHolder searchUsageHolder = new UsageService().getSearchUsageHolder();
155+
try (XContentParser jsonParser = createParser(JsonXContent.jsonXContent, restContent)) {
156+
SearchSourceBuilder source = new SearchSourceBuilder().parseXContent(jsonParser, true, searchUsageHolder, nf -> true);
157+
assertThat(source.retriever(), instanceOf(PinnedRetrieverBuilder.class));
158+
PinnedRetrieverBuilder parsed = (PinnedRetrieverBuilder) source.retriever();
159+
assertThat(parsed.retrieverName(), equalTo("my_pinned_retriever"));
160+
try (XContentParser parseSerialized = createParser(JsonXContent.jsonXContent, Strings.toString(source))) {
161+
SearchSourceBuilder deserializedSource = new SearchSourceBuilder().parseXContent(
162+
parseSerialized,
163+
true,
164+
searchUsageHolder,
165+
nf -> true
166+
);
167+
assertThat(deserializedSource.retriever(), instanceOf(PinnedRetrieverBuilder.class));
168+
PinnedRetrieverBuilder deserialized = (PinnedRetrieverBuilder) deserializedSource.retriever();
169+
assertThat(parsed, equalTo(deserialized));
170+
}
171+
}
172+
}
173+
}

0 commit comments

Comments
 (0)