Skip to content

Commit 796d94e

Browse files
jonathan-buttnerjonathan-buttner
committed
Adding integration test
1 parent 756882b commit 796d94e

File tree

1 file changed

+43
-63
lines changed
  • x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration

1 file changed

+43
-63
lines changed

‎x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/InferenceIndicesIT.java‎

Lines changed: 43 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import org.elasticsearch.action.search.SearchPhaseExecutionException;
1212
import org.elasticsearch.common.bytes.BytesReference;
1313
import org.elasticsearch.common.settings.Settings;
14+
import org.elasticsearch.core.TimeValue;
1415
import org.elasticsearch.inference.InferenceServiceExtension;
1516
import org.elasticsearch.inference.TaskType;
1617
import org.elasticsearch.license.LicenseSettings;
@@ -22,7 +23,10 @@
2223
import org.elasticsearch.xcontent.XContentFactory;
2324
import org.elasticsearch.xcontent.XContentType;
2425
import org.elasticsearch.xpack.core.LocalStateCompositeXPackPlugin;
26+
import org.elasticsearch.xpack.core.inference.InferenceContext;
2527
import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction;
28+
import org.elasticsearch.xpack.core.inference.action.InferenceAction;
29+
import org.elasticsearch.xpack.core.inference.action.InferenceActionProxy;
2630
import org.elasticsearch.xpack.core.inference.action.PutInferenceModelAction;
2731
import org.elasticsearch.xpack.core.ssl.SSLService;
2832
import org.elasticsearch.xpack.inference.InferenceIndex;
@@ -31,12 +35,10 @@
3135
import org.elasticsearch.xpack.inference.mock.TestDenseInferenceServiceExtension;
3236
import org.elasticsearch.xpack.inference.mock.TestInferenceServicePlugin;
3337
import org.elasticsearch.xpack.inference.mock.TestSparseInferenceServiceExtension;
34-
import org.junit.After;
3538

3639
import java.io.IOException;
3740
import java.nio.file.Path;
3841
import java.util.Collection;
39-
import java.util.HashMap;
4042
import java.util.List;
4143
import java.util.Map;
4244

@@ -50,7 +52,7 @@ public class InferenceIndicesIT extends ESIntegTestCase {
5052
private static final String CONFIG_ROUTER = "config";
5153
private static final String SECRETS_ROUTER = "secrets";
5254

53-
private static final Map<String, Object> BBQ_COMPATIBLE_SERVICE_SETTINGS = Map.of(
55+
private static final Map<String, Object> TEST_SERVICE_SETTINGS = Map.of(
5456
"model",
5557
"my_model",
5658
"dimensions",
@@ -61,8 +63,6 @@ public class InferenceIndicesIT extends ESIntegTestCase {
6163
"my_api_key"
6264
);
6365

64-
private final Map<String, TaskType> inferenceIds = new HashMap<>();
65-
6666
public static class LocalStateIndexSettingsInferencePlugin extends LocalStateCompositeXPackPlugin {
6767
private final InferencePlugin inferencePlugin;
6868

@@ -107,20 +107,6 @@ public Settings getSecretsIndexSettings() {
107107

108108
}
109109

110-
@After
111-
public void cleanUp() {
112-
// for (var entry : inferenceIds.entrySet()) {
113-
// assertAcked(
114-
// safeGet(
115-
// client().execute(
116-
// DeleteInferenceEndpointAction.INSTANCE,
117-
// new DeleteInferenceEndpointAction.Request(entry.getKey(), entry.getValue(), true, false)
118-
// )
119-
// )
120-
// );
121-
// }
122-
}
123-
124110
@Override
125111
protected Settings nodeSettings(int nodeOrdinal, Settings otherSettings) {
126112
return Settings.builder().put(LicenseSettings.SELF_GENERATED_LICENSE_TYPE.getKey(), "trial").build();
@@ -132,22 +118,18 @@ protected Collection<Class<? extends Plugin>> nodePlugins() {
132118
}
133119

134120
public void testRetrievingInferenceEndpoint_ThrowsException_WhenIndexNodeIsNotAvailable() throws Exception {
135-
final Settings configIndexNodeAttributes = Settings.builder().put(INDEX_ROUTER_ATTRIBUTE, CONFIG_ROUTER).build();
121+
final var configIndexNodeAttributes = Settings.builder().put(INDEX_ROUTER_ATTRIBUTE, CONFIG_ROUTER).build();
136122

137123
internalCluster().startMasterOnlyNode(configIndexNodeAttributes);
138-
final String configIndexDataNodes = internalCluster().startDataOnlyNode(configIndexNodeAttributes);
124+
final var configIndexDataNodes = internalCluster().startDataOnlyNode(configIndexNodeAttributes);
139125

140126
internalCluster().startDataOnlyNode(Settings.builder().put(INDEX_ROUTER_ATTRIBUTE, SECRETS_ROUTER).build());
141127

142-
final String inferenceId = "test-index-id";
143-
createInferenceEndpoint(TaskType.TEXT_EMBEDDING, inferenceId, BBQ_COMPATIBLE_SERVICE_SETTINGS);
128+
final var inferenceId = "test-index-id";
129+
createInferenceEndpoint(TaskType.TEXT_EMBEDDING, inferenceId, TEST_SERVICE_SETTINGS);
144130

145131
// Ensure the inference indices are created and we can retrieve the inference endpoint
146-
GetInferenceModelAction.Request getInferenceEndpointRequest = new GetInferenceModelAction.Request(
147-
inferenceId,
148-
TaskType.TEXT_EMBEDDING,
149-
true
150-
);
132+
var getInferenceEndpointRequest = new GetInferenceModelAction.Request(inferenceId, TaskType.TEXT_EMBEDDING, true);
151133
var responseFuture = client().execute(GetInferenceModelAction.INSTANCE, getInferenceEndpointRequest);
152134
assertThat(responseFuture.actionGet(TEST_REQUEST_TIMEOUT).getEndpoints().get(0).getInferenceEntityId(), equalTo(inferenceId));
153135

@@ -163,79 +145,77 @@ public void testRetrievingInferenceEndpoint_ThrowsException_WhenIndexNodeIsNotAv
163145
}
164146

165147
public void testRetrievingInferenceEndpoint_ThrowsException_WhenSecretsIndexNodeIsNotAvailable() throws Exception {
166-
final Settings configIndexNodeAttributes = Settings.builder().put(INDEX_ROUTER_ATTRIBUTE, CONFIG_ROUTER).build();
148+
final var configIndexNodeAttributes = Settings.builder().put(INDEX_ROUTER_ATTRIBUTE, CONFIG_ROUTER).build();
167149
internalCluster().startMasterOnlyNode(configIndexNodeAttributes);
168150
internalCluster().startDataOnlyNode(configIndexNodeAttributes);
169151

170-
String secretIndexDataNodes = internalCluster().startDataOnlyNode(
152+
var secretIndexDataNodes = internalCluster().startDataOnlyNode(
171153
Settings.builder().put(INDEX_ROUTER_ATTRIBUTE, SECRETS_ROUTER).build()
172154
);
173155

174-
final String inferenceId = "test-secrets-index-id";
175-
createInferenceEndpoint(TaskType.TEXT_EMBEDDING, inferenceId, BBQ_COMPATIBLE_SERVICE_SETTINGS);
156+
final var inferenceId = "test-secrets-index-id";
157+
createInferenceEndpoint(TaskType.TEXT_EMBEDDING, inferenceId, TEST_SERVICE_SETTINGS);
176158

177159
// Ensure the inference indices are created and we can retrieve the inference endpoint
178-
GetInferenceModelAction.Request getInferenceEndpointRequest = new GetInferenceModelAction.Request(
179-
inferenceId,
180-
TaskType.TEXT_EMBEDDING,
181-
true
182-
);
160+
var getInferenceEndpointRequest = new GetInferenceModelAction.Request(inferenceId, TaskType.TEXT_EMBEDDING, true);
183161
var responseFuture = client().execute(GetInferenceModelAction.INSTANCE, getInferenceEndpointRequest);
184162
assertThat(responseFuture.actionGet(TEST_REQUEST_TIMEOUT).getEndpoints().get(0).getInferenceEntityId(), equalTo(inferenceId));
185163

186164
// stop the node that holds the inference secrets index
187165
internalCluster().stopNode(secretIndexDataNodes);
188166

189-
// We should not be able to create a new inference endpoint because the secrets index is not available
190-
final String inferenceIdFailing = "test-secrets-index-id2";
191-
var responseFailureFuture = createInferenceEndpointAsync(
167+
var proxyResponse = sendInferenceProxyRequest(inferenceId);
168+
var exception = expectThrows(SearchPhaseExecutionException.class, () -> proxyResponse.actionGet(TEST_REQUEST_TIMEOUT));
169+
170+
assertThat(exception.toString(), containsString("shards failure"));
171+
assertThat(exception.toString(), containsString("Node not connected"));
172+
assertThat(exception.toString(), containsString(".secrets-inference"));
173+
}
174+
175+
private ActionFuture<InferenceAction.Response> sendInferenceProxyRequest(String inferenceId) throws IOException {
176+
final BytesReference content;
177+
try (XContentBuilder builder = XContentFactory.jsonBuilder()) {
178+
builder.startObject();
179+
builder.field("input", List.of("test input"));
180+
builder.endObject();
181+
182+
content = BytesReference.bytes(builder);
183+
}
184+
185+
var inferenceRequest = new InferenceActionProxy.Request(
192186
TaskType.TEXT_EMBEDDING,
193-
inferenceIdFailing,
194-
BBQ_COMPATIBLE_SERVICE_SETTINGS
187+
inferenceId,
188+
content,
189+
XContentType.JSON,
190+
TimeValue.THIRTY_SECONDS,
191+
false,
192+
InferenceContext.EMPTY_INSTANCE
195193
);
196-
var exception = expectThrows(SearchPhaseExecutionException.class, () -> responseFailureFuture.actionGet(TEST_REQUEST_TIMEOUT));
197194

198-
assertThat(exception.toString(), containsString("all shards failed"));
199-
assertThat(exception.toString(), containsString("Node not connected"));
200-
assertThat(exception.toString(), containsString(".inference-secrets"));
195+
return client().execute(InferenceActionProxy.INSTANCE, inferenceRequest);
201196
}
202197

203198
private void createInferenceEndpoint(TaskType taskType, String inferenceId, Map<String, Object> serviceSettings) throws IOException {
204199
var responseFuture = createInferenceEndpointAsync(taskType, inferenceId, serviceSettings);
205200
assertThat(responseFuture.actionGet(TEST_REQUEST_TIMEOUT).getModel().getInferenceEntityId(), equalTo(inferenceId));
206-
207-
inferenceIds.put(inferenceId, taskType);
208201
}
209202

210203
private ActionFuture<PutInferenceModelAction.Response> createInferenceEndpointAsync(
211204
TaskType taskType,
212205
String inferenceId,
213206
Map<String, Object> serviceSettings
214207
) throws IOException {
215-
final String service = switch (taskType) {
216-
case TEXT_EMBEDDING -> TestDenseInferenceServiceExtension.TestInferenceService.NAME;
217-
case SPARSE_EMBEDDING -> TestSparseInferenceServiceExtension.TestInferenceService.NAME;
218-
default -> throw new IllegalArgumentException("Unhandled task type [" + taskType + "]");
219-
};
220-
221208
final BytesReference content;
222209
try (XContentBuilder builder = XContentFactory.jsonBuilder()) {
223210
builder.startObject();
224-
builder.field("service", service);
211+
builder.field("service", TestDenseInferenceServiceExtension.TestInferenceService.NAME);
225212
builder.field("service_settings", serviceSettings);
226213
builder.endObject();
227214

228215
content = BytesReference.bytes(builder);
229216
}
230217

231-
PutInferenceModelAction.Request request = new PutInferenceModelAction.Request(
232-
taskType,
233-
inferenceId,
234-
content,
235-
XContentType.JSON,
236-
TEST_REQUEST_TIMEOUT
237-
);
238-
218+
var request = new PutInferenceModelAction.Request(taskType, inferenceId, content, XContentType.JSON, TEST_REQUEST_TIMEOUT);
239219
return client().execute(PutInferenceModelAction.INSTANCE, request);
240220
}
241221
}

0 commit comments

Comments
 (0)