Skip to content

Commit 6a095d2

Browse files
jonathan-buttnerjonathan-buttner
committed
Adding crud tests
1 parent d971c7c commit 6a095d2

File tree

7 files changed

+261
-26
lines changed

7 files changed

+261
-26
lines changed

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/action/PutCCMConfigurationAction.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,14 @@ public static Request parseRequest(TimeValue masterNodeTimeout, TimeValue ackTim
4949
return builder.build(masterNodeTimeout, ackTimeout);
5050
}
5151

52+
public static Request createEnabled(SecureString apiKey, TimeValue masterNodeTimeout, TimeValue ackTimeout) {
53+
return new Request(Objects.requireNonNull(apiKey), null, masterNodeTimeout, ackTimeout);
54+
}
55+
56+
public static Request createDisabled(TimeValue masterNodeTimeout, TimeValue ackTimeout) {
57+
return new Request(null, Boolean.FALSE, masterNodeTimeout, ackTimeout);
58+
}
59+
5260
private final SecureString apiKey;
5361
private final Boolean enabled;
5462

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

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
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+
* this file has been contributed to by a Generative AI
8+
*/
9+
10+
package org.elasticsearch.xpack.inference;
11+
12+
import org.elasticsearch.client.Request;
13+
import org.elasticsearch.client.Response;
14+
import org.elasticsearch.client.ResponseException;
15+
import org.elasticsearch.common.Strings;
16+
import org.elasticsearch.common.settings.SecureString;
17+
import org.elasticsearch.common.settings.Settings;
18+
import org.elasticsearch.common.util.concurrent.ThreadContext;
19+
import org.elasticsearch.core.TimeValue;
20+
import org.elasticsearch.test.cluster.ElasticsearchCluster;
21+
import org.elasticsearch.test.cluster.local.distribution.DistributionType;
22+
import org.elasticsearch.test.rest.ESRestTestCase;
23+
import org.elasticsearch.xcontent.XContentFactory;
24+
import org.elasticsearch.xcontent.XContentType;
25+
import org.elasticsearch.xpack.core.inference.action.CCMEnabledActionResponse;
26+
import org.elasticsearch.xpack.core.inference.action.PutCCMConfigurationAction;
27+
import org.junit.After;
28+
import org.junit.ClassRule;
29+
30+
import java.io.IOException;
31+
32+
import static org.elasticsearch.xpack.core.inference.action.CCMEnabledActionResponse.ENABLED_FIELD_NAME;
33+
import static org.elasticsearch.xpack.inference.InferenceBaseRestTest.assertStatusOkOrCreated;
34+
import static org.elasticsearch.xpack.inference.rest.Paths.INFERENCE_CCM_PATH;
35+
import static org.hamcrest.Matchers.containsString;
36+
import static org.hamcrest.Matchers.instanceOf;
37+
38+
public class CCMCrudDisallowedIT extends ESRestTestCase {
39+
40+
static final String PUT_METHOD = "PUT";
41+
static final String GET_METHOD = "GET";
42+
43+
@ClassRule
44+
public static ElasticsearchCluster cluster = ElasticsearchCluster.local()
45+
.distribution(DistributionType.DEFAULT)
46+
.setting("xpack.license.self_generated.type", "basic")
47+
.setting("xpack.security.enabled", "true")
48+
.setting("xpack.inference.elastic.allow_configuring_ccm", "false")
49+
// This plugin is located in the inference/qa/test-service-plugin package, look for TestInferenceServicePlugin
50+
.plugin("inference-service-test")
51+
.user("x_pack_rest_user", "x-pack-test-password")
52+
.build();
53+
54+
@Override
55+
protected String getTestRestCluster() {
56+
return cluster.getHttpAddresses();
57+
}
58+
59+
@Override
60+
protected Settings restClientSettings() {
61+
String token = basicAuthHeaderValue("x_pack_rest_user", new SecureString("x-pack-test-password".toCharArray()));
62+
return Settings.builder().put(ThreadContext.PREFIX + ".Authorization", token).build();
63+
}
64+
65+
// TODO always register the ccm endpoints but throw in the transport action logic saying that they are disabled
66+
67+
@After
68+
public void cleanup() throws IOException {
69+
// Disable CCM after each test to ensure a clean state
70+
putCCMConfiguration(PutCCMConfigurationAction.Request.createDisabled(TimeValue.THIRTY_SECONDS, TimeValue.THIRTY_SECONDS));
71+
}
72+
73+
public void testEnablesCCM_Succeeds() throws IOException {
74+
var response = putCCMConfiguration(
75+
PutCCMConfigurationAction.Request.createEnabled(
76+
new SecureString("key".toCharArray()),
77+
TimeValue.THIRTY_SECONDS,
78+
TimeValue.THIRTY_SECONDS
79+
)
80+
);
81+
82+
assertTrue(response.isEnabled());
83+
}
84+
85+
private static CCMEnabledActionResponse putCCMConfiguration(PutCCMConfigurationAction.Request request) throws IOException {
86+
var response = putRawRequest(INFERENCE_CCM_PATH, request);
87+
assertStatusOkOrCreated(response);
88+
var responseAsMap = entityAsMap(response);
89+
90+
var enabled = responseAsMap.get(ENABLED_FIELD_NAME);
91+
92+
assertThat(enabled, instanceOf(Boolean.class));
93+
return new CCMEnabledActionResponse((Boolean) enabled);
94+
}
95+
96+
private static Response putRawRequest(String endpoint, PutCCMConfigurationAction.Request actionRequest) throws IOException {
97+
var builder = XContentFactory.contentBuilder(XContentType.JSON);
98+
actionRequest.toXContent(builder, null);
99+
100+
var request = new Request(PUT_METHOD, endpoint);
101+
request.setJsonEntity(Strings.toString(builder));
102+
return client().performRequest(request);
103+
}
104+
105+
public void testEnableCCM_WithNullApiKey_NullEnabled_ThrowsException() {
106+
var request = new PutCCMConfigurationAction.Request(null, null, TimeValue.THIRTY_SECONDS, TimeValue.THIRTY_SECONDS);
107+
108+
var exception = expectThrows(IOException.class, () -> putCCMConfiguration(request));
109+
assertThat(exception.getMessage(), containsString("At least one of [api_key] or [enabled] must be provided"));
110+
}
111+
112+
public void testEnableCCM_WithoutBody_ThrowsException() {
113+
var request = new Request(PUT_METHOD, INFERENCE_CCM_PATH);
114+
115+
var exception = expectThrows(ResponseException.class, () -> client().performRequest(request));
116+
assertThat(exception.getMessage(), containsString("The body must be specified"));
117+
}
118+
119+
public void testDisableCCM_WithEnabledTrue_ThrowsException() {
120+
var request = new PutCCMConfigurationAction.Request(null, true, TimeValue.THIRTY_SECONDS, TimeValue.THIRTY_SECONDS);
121+
122+
var exception = expectThrows(IOException.class, () -> putCCMConfiguration(request));
123+
assertThat(exception.getMessage(), containsString("The [enabled] field must be set to [false] when disabling CCM"));
124+
}
125+
126+
public void testDisableCCM_Succeeds() throws IOException {
127+
var response = putCCMConfiguration(PutCCMConfigurationAction.Request.createDisabled(TimeValue.THIRTY_SECONDS, TimeValue.THIRTY_SECONDS));
128+
129+
assertFalse(response.isEnabled());
130+
}
131+
132+
public void testGetCCMConfiguration_WhenCCMDisabled_ReturnsDisabled() throws IOException {
133+
assertFalse(getCCMConfiguration().isEnabled());
134+
}
135+
136+
private CCMEnabledActionResponse getCCMConfiguration() throws IOException {
137+
var getRequest = new Request(GET_METHOD, INFERENCE_CCM_PATH);
138+
var getResponse = client().performRequest(getRequest);
139+
assertStatusOkOrCreated(getResponse);
140+
var responseAsMap = entityAsMap(getResponse);
141+
142+
var enabled = responseAsMap.get(ENABLED_FIELD_NAME);
143+
assertThat(enabled, instanceOf(Boolean.class));
144+
return new CCMEnabledActionResponse((Boolean) enabled);
145+
}
146+
147+
public void testEnablesCCM_ThenDisable() throws IOException {
148+
var response = putCCMConfiguration(
149+
PutCCMConfigurationAction.Request.createEnabled(
150+
new SecureString("key".toCharArray()),
151+
TimeValue.THIRTY_SECONDS,
152+
TimeValue.THIRTY_SECONDS
153+
)
154+
);
155+
156+
assertTrue(response.isEnabled());
157+
assertTrue(getCCMConfiguration().isEnabled());
158+
159+
response = putCCMConfiguration(PutCCMConfigurationAction.Request.createDisabled(TimeValue.THIRTY_SECONDS, TimeValue.THIRTY_SECONDS));
160+
161+
assertFalse(response.isEnabled());
162+
assertFalse(getCCMConfiguration().isEnabled());
163+
}
164+
}

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

Lines changed: 74 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111

1212
import org.elasticsearch.client.Request;
1313
import org.elasticsearch.client.Response;
14+
import org.elasticsearch.client.ResponseException;
1415
import org.elasticsearch.common.Strings;
1516
import org.elasticsearch.common.settings.SecureString;
1617
import org.elasticsearch.common.settings.Settings;
@@ -23,17 +24,22 @@
2324
import org.elasticsearch.xcontent.XContentType;
2425
import org.elasticsearch.xpack.core.inference.action.CCMEnabledActionResponse;
2526
import org.elasticsearch.xpack.core.inference.action.PutCCMConfigurationAction;
27+
import org.junit.After;
2628
import org.junit.ClassRule;
2729

2830
import java.io.IOException;
2931

3032
import static org.elasticsearch.xpack.core.inference.action.CCMEnabledActionResponse.ENABLED_FIELD_NAME;
3133
import static org.elasticsearch.xpack.inference.InferenceBaseRestTest.assertStatusOkOrCreated;
34+
import static org.elasticsearch.xpack.inference.rest.Paths.INFERENCE_CCM_PATH;
3235
import static org.hamcrest.Matchers.containsString;
3336
import static org.hamcrest.Matchers.instanceOf;
3437

3538
public class CCMCrudIT extends ESRestTestCase {
3639

40+
static final String PUT_METHOD = "PUT";
41+
static final String GET_METHOD = "GET";
42+
3743
@ClassRule
3844
public static ElasticsearchCluster cluster = ElasticsearchCluster.local()
3945
.distribution(DistributionType.DEFAULT)
@@ -56,11 +62,16 @@ protected Settings restClientSettings() {
5662
return Settings.builder().put(ThreadContext.PREFIX + ".Authorization", token).build();
5763
}
5864

59-
public void testEnablesCCM() throws IOException {
65+
@After
66+
public void cleanup() throws IOException {
67+
// Disable CCM after each test to ensure a clean state
68+
putCCMConfiguration(PutCCMConfigurationAction.Request.createDisabled(TimeValue.THIRTY_SECONDS, TimeValue.THIRTY_SECONDS));
69+
}
70+
71+
public void testEnablesCCM_Succeeds() throws IOException {
6072
var response = putCCMConfiguration(
61-
new PutCCMConfigurationAction.Request(
73+
PutCCMConfigurationAction.Request.createEnabled(
6274
new SecureString("key".toCharArray()),
63-
null,
6475
TimeValue.THIRTY_SECONDS,
6576
TimeValue.THIRTY_SECONDS
6677
)
@@ -70,7 +81,7 @@ public void testEnablesCCM() throws IOException {
7081
}
7182

7283
private static CCMEnabledActionResponse putCCMConfiguration(PutCCMConfigurationAction.Request request) throws IOException {
73-
var response = putRawRequest("_inference/_ccm", request);
84+
var response = putRawRequest(INFERENCE_CCM_PATH, request);
7485
assertStatusOkOrCreated(response);
7586
var responseAsMap = entityAsMap(response);
7687

@@ -80,25 +91,72 @@ private static CCMEnabledActionResponse putCCMConfiguration(PutCCMConfigurationA
8091
return new CCMEnabledActionResponse((Boolean) enabled);
8192
}
8293

94+
private static Response putRawRequest(String endpoint, PutCCMConfigurationAction.Request actionRequest) throws IOException {
95+
var builder = XContentFactory.contentBuilder(XContentType.JSON);
96+
actionRequest.toXContent(builder, null);
97+
98+
var request = new Request(PUT_METHOD, endpoint);
99+
request.setJsonEntity(Strings.toString(builder));
100+
return client().performRequest(request);
101+
}
102+
83103
public void testEnableCCM_WithNullApiKey_NullEnabled_ThrowsException() {
84-
var request = new PutCCMConfigurationAction.Request(
85-
null,
86-
null,
87-
TimeValue.THIRTY_SECONDS,
88-
TimeValue.THIRTY_SECONDS
89-
);
104+
var request = new PutCCMConfigurationAction.Request(null, null, TimeValue.THIRTY_SECONDS, TimeValue.THIRTY_SECONDS);
90105

91106
var exception = expectThrows(IOException.class, () -> putCCMConfiguration(request));
92107
assertThat(exception.getMessage(), containsString("At least one of [api_key] or [enabled] must be provided"));
93108
}
94109

95-
private static Response putRawRequest(String endpoint, PutCCMConfigurationAction.Request actionRequest) throws IOException {
96-
var builder = XContentFactory.contentBuilder(XContentType.JSON);
97-
actionRequest.toXContent(builder, null);
110+
public void testEnableCCM_WithoutBody_ThrowsException() {
111+
var request = new Request(PUT_METHOD, INFERENCE_CCM_PATH);
98112

99-
var request = new Request("PUT", endpoint);
100-
request.setJsonEntity(Strings.toString(builder));
101-
return client().performRequest(request);
113+
var exception = expectThrows(ResponseException.class, () -> client().performRequest(request));
114+
assertThat(exception.getMessage(), containsString("The body must be specified"));
102115
}
103116

117+
public void testDisableCCM_WithEnabledTrue_ThrowsException() {
118+
var request = new PutCCMConfigurationAction.Request(null, true, TimeValue.THIRTY_SECONDS, TimeValue.THIRTY_SECONDS);
119+
120+
var exception = expectThrows(IOException.class, () -> putCCMConfiguration(request));
121+
assertThat(exception.getMessage(), containsString("The [enabled] field must be set to [false] when disabling CCM"));
122+
}
123+
124+
public void testDisableCCM_Succeeds() throws IOException {
125+
var response = putCCMConfiguration(PutCCMConfigurationAction.Request.createDisabled(TimeValue.THIRTY_SECONDS, TimeValue.THIRTY_SECONDS));
126+
127+
assertFalse(response.isEnabled());
128+
}
129+
130+
public void testGetCCMConfiguration_WhenCCMDisabled_ReturnsDisabled() throws IOException {
131+
assertFalse(getCCMConfiguration().isEnabled());
132+
}
133+
134+
private CCMEnabledActionResponse getCCMConfiguration() throws IOException {
135+
var getRequest = new Request(GET_METHOD, INFERENCE_CCM_PATH);
136+
var getResponse = client().performRequest(getRequest);
137+
assertStatusOkOrCreated(getResponse);
138+
var responseAsMap = entityAsMap(getResponse);
139+
140+
var enabled = responseAsMap.get(ENABLED_FIELD_NAME);
141+
assertThat(enabled, instanceOf(Boolean.class));
142+
return new CCMEnabledActionResponse((Boolean) enabled);
143+
}
144+
145+
public void testEnablesCCM_ThenDisable() throws IOException {
146+
var response = putCCMConfiguration(
147+
PutCCMConfigurationAction.Request.createEnabled(
148+
new SecureString("key".toCharArray()),
149+
TimeValue.THIRTY_SECONDS,
150+
TimeValue.THIRTY_SECONDS
151+
)
152+
);
153+
154+
assertTrue(response.isEnabled());
155+
assertTrue(getCCMConfiguration().isEnabled());
156+
157+
response = putCCMConfiguration(PutCCMConfigurationAction.Request.createDisabled(TimeValue.THIRTY_SECONDS, TimeValue.THIRTY_SECONDS));
158+
159+
assertFalse(response.isEnabled());
160+
assertFalse(getCCMConfiguration().isEnabled());
161+
}
104162
}

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@
6565
import org.elasticsearch.xpack.core.inference.action.GetRerankerWindowSizeAction;
6666
import org.elasticsearch.xpack.core.inference.action.InferenceAction;
6767
import org.elasticsearch.xpack.core.inference.action.InferenceActionProxy;
68+
import org.elasticsearch.xpack.core.inference.action.PutCCMConfigurationAction;
6869
import org.elasticsearch.xpack.core.inference.action.PutInferenceModelAction;
6970
import org.elasticsearch.xpack.core.inference.action.UnifiedCompletionAction;
7071
import org.elasticsearch.xpack.core.inference.action.UpdateInferenceModelAction;
@@ -78,6 +79,7 @@
7879
import org.elasticsearch.xpack.inference.action.TransportInferenceAction;
7980
import org.elasticsearch.xpack.inference.action.TransportInferenceActionProxy;
8081
import org.elasticsearch.xpack.inference.action.TransportInferenceUsageAction;
82+
import org.elasticsearch.xpack.inference.action.TransportPutCCMConfigurationAction;
8183
import org.elasticsearch.xpack.inference.action.TransportPutInferenceModelAction;
8284
import org.elasticsearch.xpack.inference.action.TransportUnifiedCompletionInferenceAction;
8385
import org.elasticsearch.xpack.inference.action.TransportUpdateInferenceModelAction;
@@ -116,6 +118,7 @@
116118
import org.elasticsearch.xpack.inference.rest.RestGetInferenceModelAction;
117119
import org.elasticsearch.xpack.inference.rest.RestGetInferenceServicesAction;
118120
import org.elasticsearch.xpack.inference.rest.RestInferenceAction;
121+
import org.elasticsearch.xpack.inference.rest.RestPutCCMConfigurationAction;
119122
import org.elasticsearch.xpack.inference.rest.RestPutInferenceModelAction;
120123
import org.elasticsearch.xpack.inference.rest.RestStreamInferenceAction;
121124
import org.elasticsearch.xpack.inference.rest.RestUpdateInferenceModelAction;
@@ -244,7 +247,10 @@ public InferencePlugin(Settings settings) {
244247
public List<ActionHandler> getActions() {
245248
CCMSettings.ALLOW_CONFIGURING_CCM.get(settings);
246249
List<ActionHandler> ccmActions = ccmFeature.get().allowConfiguringCcm()
247-
? List.of(new ActionHandler(GetCCMConfigurationAction.INSTANCE, TransportGetCCMConfigurationAction.class))
250+
? List.of(
251+
new ActionHandler(GetCCMConfigurationAction.INSTANCE, TransportGetCCMConfigurationAction.class),
252+
new ActionHandler(PutCCMConfigurationAction.INSTANCE, TransportPutCCMConfigurationAction.class)
253+
)
248254
: List.of();
249255

250256
return Stream.of(
@@ -278,7 +284,9 @@ public List<RestHandler> getRestHandlers(
278284
Supplier<DiscoveryNodes> nodesInCluster,
279285
Predicate<NodeFeature> clusterSupportsFeature
280286
) {
281-
List<RestHandler> ccmHandler = ccmFeature.get().allowConfiguringCcm() ? List.of(new RestGetCCMConfigurationAction()) : List.of();
287+
List<RestHandler> ccmHandler = ccmFeature.get().allowConfiguringCcm()
288+
? List.of(new RestGetCCMConfigurationAction(), new RestPutCCMConfigurationAction())
289+
: List.of();
282290
List<RestHandler> handlers = List.of(
283291
new RestInferenceAction(),
284292
new RestStreamInferenceAction(threadPoolSetOnce),

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/TransportPutCCMConfigurationAction.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
import org.elasticsearch.threadpool.ThreadPool;
2424
import org.elasticsearch.transport.TransportService;
2525
import org.elasticsearch.xpack.core.inference.action.CCMEnabledActionResponse;
26-
import org.elasticsearch.xpack.core.inference.action.GetCCMConfigurationAction;
2726
import org.elasticsearch.xpack.core.inference.action.PutCCMConfigurationAction;
2827
import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMModel;
2928
import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMService;
@@ -45,7 +44,7 @@ public TransportPutCCMConfigurationAction(
4544
ProjectResolver projectResolver
4645
) {
4746
super(
48-
GetCCMConfigurationAction.NAME,
47+
PutCCMConfigurationAction.NAME,
4948
transportService,
5049
clusterService,
5150
threadPool,

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rest/Paths.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ public final class Paths {
1212
static final String INFERENCE_ID = "inference_id";
1313
static final String TASK_TYPE_OR_INFERENCE_ID = "task_type_or_id";
1414
static final String TASK_TYPE = "task_type";
15-
static final String INFERENCE_CCM_PATH = "_inference/_ccm";
15+
public static final String INFERENCE_CCM_PATH = "_inference/_ccm";
1616
static final String INFERENCE_ID_PATH = "_inference/{" + TASK_TYPE_OR_INFERENCE_ID + "}";
1717
static final String TASK_TYPE_INFERENCE_ID_PATH = "_inference/{" + TASK_TYPE_OR_INFERENCE_ID + "}/{" + INFERENCE_ID + "}";
1818
static final String INFERENCE_DIAGNOSTICS_PATH = "_inference/.diagnostics";

0 commit comments

Comments
 (0)