Skip to content

Commit 8f374e1

Browse files
Creating more tests and service class
1 parent 2c63c9a commit 8f374e1

File tree

13 files changed

+723
-47
lines changed

13 files changed

+723
-47
lines changed
Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
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.core.inference.action;
9+
10+
import org.elasticsearch.action.ActionRequest;
11+
import org.elasticsearch.action.ActionRequestValidationException;
12+
import org.elasticsearch.action.ActionResponse;
13+
import org.elasticsearch.action.ActionType;
14+
import org.elasticsearch.common.io.stream.StreamInput;
15+
import org.elasticsearch.common.io.stream.StreamOutput;
16+
import org.elasticsearch.xcontent.ToXContentObject;
17+
import org.elasticsearch.xcontent.XContentBuilder;
18+
19+
import java.io.IOException;
20+
import java.util.Objects;
21+
22+
public class GetCCMConfigurationAction extends ActionType<GetCCMConfigurationAction.Response> {
23+
24+
public static final GetCCMConfigurationAction INSTANCE = new GetCCMConfigurationAction();
25+
public static final String NAME = "cluster:monitor/xpack/inference/ccm/get";
26+
27+
public GetCCMConfigurationAction() {
28+
super(NAME);
29+
}
30+
31+
public static class Request extends ActionRequest {
32+
33+
public Request() {}
34+
35+
public Request(StreamInput in) throws IOException {
36+
super(in);
37+
}
38+
39+
@Override
40+
public ActionRequestValidationException validate() {
41+
return null;
42+
}
43+
44+
@Override
45+
public boolean equals(Object o) {
46+
if (this == o) return true;
47+
if (o == null || getClass() != o.getClass()) return false;
48+
return true;
49+
}
50+
51+
@Override
52+
public int hashCode() {
53+
// The class doesn't have any members at the moment so return the same hash code
54+
return Objects.hash(NAME);
55+
}
56+
}
57+
58+
public static class Response extends ActionResponse implements ToXContentObject {
59+
private final boolean enabled;
60+
61+
public Response(boolean enabled) {
62+
this.enabled = enabled;
63+
}
64+
65+
public Response(StreamInput in) throws IOException {
66+
this.enabled = in.readBoolean();
67+
}
68+
69+
@Override
70+
public void writeTo(StreamOutput out) throws IOException {
71+
out.writeBoolean(enabled);
72+
}
73+
74+
public boolean isEnabled() {
75+
return enabled;
76+
}
77+
78+
@Override
79+
public boolean equals(Object o) {
80+
if (o == null || getClass() != o.getClass()) return false;
81+
Response response = (Response) o;
82+
return enabled == response.enabled;
83+
}
84+
85+
@Override
86+
public int hashCode() {
87+
return Objects.hashCode(enabled);
88+
}
89+
90+
@Override
91+
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
92+
builder.startObject();
93+
builder.field("enabled", enabled);
94+
builder.endObject();
95+
return builder;
96+
}
97+
}
98+
}
Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
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.inference.integration;
9+
10+
import org.elasticsearch.ElasticsearchException;
11+
import org.elasticsearch.ResourceNotFoundException;
12+
import org.elasticsearch.action.ActionListener;
13+
import org.elasticsearch.action.DocWriteResponse;
14+
import org.elasticsearch.action.support.PlainActionFuture;
15+
import org.elasticsearch.action.support.WriteRequest;
16+
import org.elasticsearch.common.settings.SecureString;
17+
import org.elasticsearch.core.TimeValue;
18+
import org.elasticsearch.index.query.QueryBuilders;
19+
import org.elasticsearch.plugins.Plugin;
20+
import org.elasticsearch.reindex.ReindexPlugin;
21+
import org.elasticsearch.test.ESSingleNodeTestCase;
22+
import org.elasticsearch.xcontent.XContentType;
23+
import org.elasticsearch.xpack.inference.LocalStateInferencePlugin;
24+
import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMIndex;
25+
import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMModel;
26+
27+
import java.util.Collection;
28+
import java.util.Objects;
29+
30+
import static org.elasticsearch.xpack.inference.services.elastic.ccm.CCMPersistentStorageService.CCM_DOC_ID;
31+
import static org.hamcrest.Matchers.containsString;
32+
import static org.hamcrest.Matchers.is;
33+
34+
public abstract class BaseCCMIT extends ESSingleNodeTestCase {
35+
36+
private Provider provider;
37+
38+
public interface Provider {
39+
void store(CCMModel ccmModel, ActionListener<Void> listener);
40+
void get(ActionListener<CCMModel> listener);
41+
void delete(ActionListener<Void> listener);
42+
}
43+
44+
public BaseCCMIT(Provider provider) {
45+
this.provider = Objects.requireNonNull(provider);
46+
}
47+
48+
@Override
49+
protected Collection<Class<? extends Plugin>> getPlugins() {
50+
return pluginList(ReindexPlugin.class, LocalStateInferencePlugin.class);
51+
}
52+
53+
public void testStoreAndGetCCMModel() {
54+
assertStoreCCMConfiguration();
55+
}
56+
57+
private void assertStoreCCMConfiguration() {
58+
assertStoreCCMConfiguration("secret");
59+
}
60+
61+
private void assertStoreCCMConfiguration(String apiKey) {
62+
var ccmModel = new CCMModel(new SecureString(apiKey.toCharArray()));
63+
var storeListener = new PlainActionFuture<Void>();
64+
provider.store(ccmModel, storeListener);
65+
66+
assertNull(storeListener.actionGet(TimeValue.THIRTY_SECONDS));
67+
68+
var getListener = new PlainActionFuture<CCMModel>();
69+
provider.get(getListener);
70+
71+
assertThat(getListener.actionGet(TimeValue.THIRTY_SECONDS), is(ccmModel));
72+
}
73+
74+
public void testStore_OverwritesConfiguration_WhenItAlreadyExists() {
75+
assertStoreCCMConfiguration();
76+
assertStoreCCMConfiguration("new_secret");
77+
78+
var results = client().prepareSearch(CCMIndex.INDEX_PATTERN)
79+
.setQuery(QueryBuilders.idsQuery().addIds(CCM_DOC_ID))
80+
.execute().actionGet(TimeValue.THIRTY_SECONDS);
81+
82+
assertThat(results.getHits().getHits().length, is(1));
83+
}
84+
85+
public void testGet_ThrowsResourceNotFoundException_WhenCCMIndexDoesNotExist() {
86+
assertCCMResourceDoesNotExist();
87+
}
88+
89+
private void assertCCMResourceDoesNotExist() {
90+
var getListener = new PlainActionFuture<CCMModel>();
91+
provider.get(getListener);
92+
93+
var exception = expectThrows(ResourceNotFoundException.class, () -> getListener.actionGet(TimeValue.THIRTY_SECONDS));
94+
assertThat(exception.getMessage(), is("CCM configuration not found"));
95+
}
96+
97+
public void testGet_ThrowsResourceNotFoundException_WhenCCMConfigurationDocumentDoesNotExist() {
98+
storeCorruptCCMModel("id");
99+
100+
assertCCMResourceDoesNotExist();
101+
}
102+
103+
public void testGetCCMModel_ThrowsException_WhenStoredModelIsCorrupted() {
104+
storeCorruptCCMModel(CCM_DOC_ID);
105+
106+
var getListener = new PlainActionFuture<CCMModel>();
107+
provider.get(getListener);
108+
109+
var exception = expectThrows(ElasticsearchException.class, () -> getListener.actionGet(TimeValue.THIRTY_SECONDS));
110+
assertThat(exception.getMessage(), containsString("Failed to retrieve CCM configuration"));
111+
assertThat(exception.getCause().getMessage(), containsString("Required [api_key]"));
112+
}
113+
114+
private void storeCorruptCCMModel(String id) {
115+
var corruptedSource = """
116+
{
117+
118+
}
119+
""";
120+
121+
var response = client().prepareIndex()
122+
.setSource(corruptedSource, XContentType.JSON)
123+
.setIndex(CCMIndex.INDEX_NAME)
124+
.setId(id)
125+
.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
126+
.execute()
127+
.actionGet(TimeValue.THIRTY_SECONDS);
128+
129+
assertThat(response.getResult(), is(DocWriteResponse.Result.CREATED));
130+
}
131+
132+
public void testDelete_DoesNotThrow_WhenTheConfigurationDoesNotExist() {
133+
var listener = new PlainActionFuture<Void>();
134+
provider.delete(listener);
135+
136+
assertNull(listener.actionGet(TimeValue.THIRTY_SECONDS));
137+
assertCCMResourceDoesNotExist();
138+
}
139+
140+
public void testDelete_RemovesCCMConfiguration() {
141+
assertStoreCCMConfiguration();
142+
143+
var listener = new PlainActionFuture<Void>();
144+
provider.delete(listener);
145+
146+
assertNull(listener.actionGet(TimeValue.THIRTY_SECONDS));
147+
assertCCMResourceDoesNotExist();
148+
}
149+
}

0 commit comments

Comments
 (0)