Skip to content

Commit f6b2d74

Browse files
jonathan-buttnerjonathan-buttner
committed
Working integration tests
1 parent 839ebe0 commit f6b2d74

File tree

8 files changed

+116
-232
lines changed

8 files changed

+116
-232
lines changed

‎x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetModelsWithElasticInferenceServiceIT.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import static org.elasticsearch.xpack.inference.InferenceBaseRestTest.getAllModels;
2020
import static org.elasticsearch.xpack.inference.InferenceBaseRestTest.getModels;
21+
import static org.elasticsearch.xpack.inference.services.elastic.InternalPreconfiguredEndpoints.DEFAULT_ELSER_ENDPOINT_ID_V2;
2122
import static org.hamcrest.Matchers.hasSize;
2223
import static org.hamcrest.Matchers.is;
2324

@@ -51,7 +52,7 @@ public void testGetDefaultEndpoints() throws IOException {
5152

5253
assertInferenceIdTaskType(allModels, ".rainbow-sprinkles-elastic", TaskType.CHAT_COMPLETION);
5354
assertInferenceIdTaskType(allModels, ".gp-llm-v2-chat_completion", TaskType.CHAT_COMPLETION);
54-
assertInferenceIdTaskType(allModels, ".elser-2-elastic", TaskType.SPARSE_EMBEDDING);
55+
assertInferenceIdTaskType(allModels, DEFAULT_ELSER_ENDPOINT_ID_V2, TaskType.SPARSE_EMBEDDING);
5556
assertInferenceIdTaskType(allModels, ".jina-embeddings-v3", TaskType.TEXT_EMBEDDING);
5657
assertInferenceIdTaskType(allModels, ".elastic-rerank-v1", TaskType.RERANK);
5758
}

‎x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/MockElasticInferenceServiceAuthorizationServer.java‎

Lines changed: 69 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -23,38 +23,79 @@ public class MockElasticInferenceServiceAuthorizationServer implements TestRule
2323
private static final Logger logger = LogManager.getLogger(MockElasticInferenceServiceAuthorizationServer.class);
2424
private final MockWebServer webServer = new MockWebServer();
2525

26-
public static MockElasticInferenceServiceAuthorizationServer enabledWithRainbowSprinklesAndElser() {
27-
var server = new MockElasticInferenceServiceAuthorizationServer();
28-
29-
server.enqueueAuthorizeAllModelsResponse();
30-
return server;
31-
}
32-
3326
public void enqueueAuthorizeAllModelsResponse() {
3427
String responseJson = """
3528
{
36-
"models": [
37-
{
38-
"model_name": "rainbow-sprinkles",
39-
"task_types": ["chat"]
40-
},
41-
{
42-
"model_name": "gp-llm-v2",
43-
"task_types": ["chat"]
44-
},
45-
{
46-
"model_name": "elser_model_2",
47-
"task_types": ["embed/text/sparse"]
48-
},
49-
{
50-
"model_name": "jina-embeddings-v3",
51-
"task_types": ["embed/text/dense"]
52-
},
53-
{
54-
"model_name": "elastic-rerank-v1",
55-
"task_types": ["rerank/text/text-similarity"]
29+
"inference_endpoints": [
30+
{
31+
"id": ".rainbow-sprinkles-elastic",
32+
"model_name": "rainbow-sprinkles",
33+
"task_type": "chat_completion",
34+
"status": "ga",
35+
"properties": [
36+
"multilingual"
37+
],
38+
"release_date": "2024-05-01",
39+
"end_of_life_date": "2025-12-31"
40+
},
41+
{
42+
"id": ".gp-llm-v2-chat_completion",
43+
"model_name": "gp-llm-v2",
44+
"task_type": "chat_completion",
45+
"status": "ga",
46+
"properties": [
47+
"multilingual"
48+
],
49+
"release_date": "2024-05-01",
50+
"end_of_life_date": "2025-12-31"
51+
},
52+
{
53+
"id": ".elser-2-elastic",
54+
"model_name": "elser_model_2",
55+
"task_type": "sparse_embedding",
56+
"status": "preview",
57+
"properties": [
58+
"english"
59+
],
60+
"release_date": "2024-05-01",
61+
"configuration": {
62+
"chunking_settings": {
63+
"strategy": "sentence",
64+
"max_chunk_size": 250,
65+
"sentence_overlap": 1
5666
}
57-
]
67+
}
68+
},
69+
{
70+
"id": ".jina-embeddings-v3",
71+
"model_name": "jina-embeddings-v3",
72+
"task_type": "text_embedding",
73+
"status": "beta",
74+
"properties": [
75+
"multilingual",
76+
"open-weights"
77+
],
78+
"release_date": "2024-05-01",
79+
"configuration": {
80+
"similarity": "cosine",
81+
"dimensions": 1024,
82+
"element_type": "float",
83+
"chunking_settings": {
84+
"strategy": "word",
85+
"max_chunk_size": 500,
86+
"overlap": 2
87+
}
88+
}
89+
},
90+
{
91+
"id": ".elastic-rerank-v1",
92+
"model_name": "elastic-rerank-v1",
93+
"task_type": "rerank",
94+
"status": "preview",
95+
"properties": [],
96+
"release_date": "2024-05-01"
97+
}
98+
]
5899
}
59100
""";
60101

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

Lines changed: 33 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import org.elasticsearch.xpack.inference.registry.ModelRegistry;
2626
import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceService;
2727
import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSettings;
28-
import org.elasticsearch.xpack.inference.services.elastic.InternalPreconfiguredEndpoints;
2928
import org.elasticsearch.xpack.inference.services.elastic.authorization.AuthorizationPoller;
3029
import org.elasticsearch.xpack.inference.services.elastic.authorization.AuthorizationTaskExecutor;
3130
import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMSettings;
@@ -37,6 +36,7 @@
3736
import java.io.IOException;
3837
import java.util.Collection;
3938
import java.util.List;
39+
import java.util.Set;
4040
import java.util.concurrent.atomic.AtomicReference;
4141
import java.util.function.Function;
4242
import java.util.stream.Collectors;
@@ -47,6 +47,30 @@
4747
import static org.hamcrest.Matchers.not;
4848

4949
public class AuthorizationTaskExecutorIT extends ESSingleNodeTestCase {
50+
51+
// rainbow-sprinkles
52+
public static final String DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1 = ".rainbow-sprinkles-elastic";
53+
54+
// gp-llm-v2
55+
public static final String GP_LLM_V2_CHAT_COMPLETION_ENDPOINT_ID = ".gp-llm-v2-chat_completion";
56+
57+
// elser-2
58+
public static final String DEFAULT_ELSER_ENDPOINT_ID_V2 = ".elser-2-elastic";
59+
60+
// multilingual-text-embed
61+
public static final String DEFAULT_MULTILINGUAL_EMBED_ENDPOINT_ID = ".jina-embeddings-v3";
62+
63+
// rerank-v1
64+
public static final String DEFAULT_RERANK_ENDPOINT_ID_V1 = ".elastic-rerank-v1";
65+
66+
public static final Set<String> EIS_PRECONFIGURED_ENDPOINT_IDS = Set.of(
67+
DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1,
68+
GP_LLM_V2_CHAT_COMPLETION_ENDPOINT_ID,
69+
DEFAULT_ELSER_ENDPOINT_ID_V2,
70+
DEFAULT_MULTILINGUAL_EMBED_ENDPOINT_ID,
71+
DEFAULT_RERANK_ENDPOINT_ID_V1
72+
);
73+
5074
public static final String AUTH_TASK_ACTION = AuthorizationPoller.TASK_NAME + "[c]";
5175

5276
public static final String EMPTY_AUTH_RESPONSE = """
@@ -94,7 +118,7 @@ public void shutdown() {
94118
static void removeEisPreconfiguredEndpoints(ModelRegistry modelRegistry) {
95119
// Delete all the eis preconfigured endpoints
96120
var listener = new PlainActionFuture<Boolean>();
97-
modelRegistry.deleteModels(InternalPreconfiguredEndpoints.EIS_PRECONFIGURED_ENDPOINT_IDS, listener);
121+
modelRegistry.deleteModels(EIS_PRECONFIGURED_ENDPOINT_IDS, listener);
98122
listener.actionGet(TimeValue.THIRTY_SECONDS);
99123
}
100124

@@ -149,7 +173,7 @@ static void assertNoAuthorizedEisEndpoints(
149173
var eisEndpoints = getEisEndpoints(modelRegistry);
150174
assertThat(eisEndpoints, empty());
151175

152-
for (String eisPreconfiguredEndpoints : InternalPreconfiguredEndpoints.EIS_PRECONFIGURED_ENDPOINT_IDS) {
176+
for (String eisPreconfiguredEndpoints : EIS_PRECONFIGURED_ENDPOINT_IDS) {
153177
assertFalse(modelRegistry.containsPreconfiguredInferenceEndpointId(eisPreconfiguredEndpoints));
154178
}
155179
}
@@ -250,15 +274,13 @@ static void assertChatCompletionEndpointExists(ModelRegistry modelRegistry) {
250274

251275
var rainbowSprinklesModel = eisEndpoints.get(0);
252276
assertChatCompletionUnparsedModel(rainbowSprinklesModel);
253-
assertTrue(
254-
modelRegistry.containsPreconfiguredInferenceEndpointId(InternalPreconfiguredEndpoints.DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1)
255-
);
277+
assertTrue(modelRegistry.containsPreconfiguredInferenceEndpointId(DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1));
256278
}
257279

258280
static void assertChatCompletionUnparsedModel(UnparsedModel rainbowSprinklesModel) {
259281
assertThat(rainbowSprinklesModel.taskType(), is(TaskType.CHAT_COMPLETION));
260282
assertThat(rainbowSprinklesModel.service(), is(ElasticInferenceService.NAME));
261-
assertThat(rainbowSprinklesModel.inferenceEntityId(), is(InternalPreconfiguredEndpoints.DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1));
283+
assertThat(rainbowSprinklesModel.inferenceEntityId(), is(DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1));
262284
}
263285

264286
public void testCreatesChatCompletion_AndThenCreatesTextEmbedding() throws Exception {
@@ -293,12 +315,12 @@ public void testCreatesChatCompletion_AndThenCreatesTextEmbedding() throws Excep
293315
var eisEndpoints = getEisEndpoints().stream().collect(Collectors.toMap(UnparsedModel::inferenceEntityId, Function.identity()));
294316
assertThat(eisEndpoints.size(), is(2));
295317

296-
assertTrue(eisEndpoints.containsKey(InternalPreconfiguredEndpoints.DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1));
297-
assertChatCompletionUnparsedModel(eisEndpoints.get(InternalPreconfiguredEndpoints.DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1));
318+
assertTrue(eisEndpoints.containsKey(DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1));
319+
assertChatCompletionUnparsedModel(eisEndpoints.get(DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1));
298320

299-
assertTrue(eisEndpoints.containsKey(InternalPreconfiguredEndpoints.DEFAULT_MULTILINGUAL_EMBED_ENDPOINT_ID));
321+
assertTrue(eisEndpoints.containsKey(DEFAULT_MULTILINGUAL_EMBED_ENDPOINT_ID));
300322

301-
var textEmbeddingEndpoint = eisEndpoints.get(InternalPreconfiguredEndpoints.DEFAULT_MULTILINGUAL_EMBED_ENDPOINT_ID);
323+
var textEmbeddingEndpoint = eisEndpoints.get(DEFAULT_MULTILINGUAL_EMBED_ENDPOINT_ID);
302324
assertThat(textEmbeddingEndpoint.taskType(), is(TaskType.TEXT_EMBEDDING));
303325
assertThat(textEmbeddingEndpoint.service(), is(ElasticInferenceService.NAME));
304326
}

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

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
import org.elasticsearch.xpack.inference.LocalStateInferencePlugin;
1919
import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceService;
2020
import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSettings;
21-
import org.elasticsearch.xpack.inference.services.elastic.InternalPreconfiguredEndpoints;
2221
import org.elasticsearch.xpack.inference.services.elastic.authorization.AuthorizationPoller;
2322
import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMSettings;
2423
import org.junit.AfterClass;
@@ -34,6 +33,7 @@
3433

3534
import static org.elasticsearch.xpack.inference.external.http.Utils.getUrl;
3635
import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.AUTHORIZED_RAINBOW_SPRINKLES_RESPONSE;
36+
import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1;
3737
import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.EMPTY_AUTH_RESPONSE;
3838
import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.cancelAuthorizationTask;
3939
import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.waitForTask;
@@ -130,10 +130,7 @@ public void testAuthorizationTaskGetsRelocatedToAnotherNode_WhenTheNodeThatIsRun
130130

131131
var rainbowSprinklesEndpoint = eisEndpoints.get(0);
132132
assertThat(rainbowSprinklesEndpoint.getService(), is(ElasticInferenceService.NAME));
133-
assertThat(
134-
rainbowSprinklesEndpoint.getInferenceEntityId(),
135-
is(InternalPreconfiguredEndpoints.DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1)
136-
);
133+
assertThat(rainbowSprinklesEndpoint.getInferenceEntityId(), is(DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1));
137134
assertThat(rainbowSprinklesEndpoint.getTaskType(), is(TaskType.CHAT_COMPLETION));
138135
});
139136
}

0 commit comments

Comments
 (0)