Skip to content

Commit 642ac18

Browse files
jonathan-buttnerjonathan-buttner
committed
Adding more tests for enablement
1 parent ee58e1b commit 642ac18

File tree

7 files changed

+258
-33
lines changed

7 files changed

+258
-33
lines changed

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

Lines changed: 16 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -462,7 +462,21 @@ private CCMRelatedComponents createCCMDependentComponents(
462462
) {
463463
var ccmEnablementService = new CCMEnablementService(services.clusterService(), services.featureService(), ccmFeature);
464464
var ccmPersistentStorageService = new CCMPersistentStorageService(services.client());
465-
var ccmService = new CCMService(ccmPersistentStorageService, ccmEnablementService, services.projectResolver(), services.client());
465+
var ccmCache = new CCMCache(
466+
ccmPersistentStorageService,
467+
services.clusterService(),
468+
settings,
469+
services.featureService(),
470+
services.projectResolver(),
471+
services.client()
472+
);
473+
var ccmService = new CCMService(
474+
ccmPersistentStorageService,
475+
ccmEnablementService,
476+
ccmCache,
477+
services.projectResolver(),
478+
services.client()
479+
);
466480
var ccmAuthApplierFactory = new CCMAuthenticationApplierFactory(ccmFeature, ccmService);
467481

468482
var authorizationHandler = new ElasticInferenceServiceAuthorizationRequestHandler(
@@ -491,20 +505,7 @@ private CCMRelatedComponents createCCMDependentComponents(
491505
authTaskExecutor.startAndLazyCreateTask();
492506

493507
return new CCMRelatedComponents(
494-
List.of(
495-
authorizationHandler,
496-
authTaskExecutor,
497-
ccmService,
498-
ccmPersistentStorageService,
499-
new CCMCache(
500-
ccmPersistentStorageService,
501-
services.clusterService(),
502-
settings,
503-
services.featureService(),
504-
services.projectResolver(),
505-
services.client()
506-
)
507-
),
508+
List.of(authorizationHandler, authTaskExecutor, ccmService, ccmPersistentStorageService, ccmCache, ccmEnablementService),
508509
ccmAuthApplierFactory
509510
);
510511
}

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ccm/CCMAuthenticationApplierFactory.java‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,10 @@
88
package org.elasticsearch.xpack.inference.services.elastic.ccm;
99

1010
import org.apache.http.client.methods.HttpRequestBase;
11+
import org.apache.logging.log4j.LogManager;
12+
import org.apache.logging.log4j.Logger;
1113
import org.elasticsearch.ElasticsearchStatusException;
14+
import org.elasticsearch.ResourceNotFoundException;
1215
import org.elasticsearch.action.ActionListener;
1316
import org.elasticsearch.action.support.SubscribableListener;
1417
import org.elasticsearch.common.settings.SecureString;
@@ -26,6 +29,7 @@
2629
public class CCMAuthenticationApplierFactory {
2730

2831
public static final NoopApplier NOOP_APPLIER = new NoopApplier();
32+
private static final Logger logger = LogManager.getLogger(CCMAuthenticationApplierFactory.class);
2933

3034
private final CCMFeature ccmFeature;
3135
private final CCMService ccmService;
@@ -56,7 +60,26 @@ public void getAuthenticationApplier(ActionListener<AuthApplier> listener) {
5660
return;
5761
}
5862

59-
ccmService.getConfiguration(ccmModelListener);
63+
var consistencyListener = ccmModelListener.delegateResponse((delegate, e) -> {
64+
if (e instanceof ResourceNotFoundException) {
65+
logger.atWarn()
66+
.withThrowable(e)
67+
.log("CCM cluster state indicates CCM is enabled but no configuration was found using the cache");
68+
listener.onFailure(
69+
new ElasticsearchStatusException(
70+
"Cloud connected mode configuration is in an inconsistent state. "
71+
+ "Please try configuring it again using PUT {}",
72+
RestStatus.BAD_REQUEST,
73+
INFERENCE_CCM_PATH
74+
)
75+
);
76+
return;
77+
}
78+
79+
delegate.onFailure(e);
80+
});
81+
82+
ccmService.getConfiguration(consistencyListener);
6083
}).<AuthApplier>andThenApply(ccmModel -> new AuthenticationHeaderApplier(ccmModel.apiKey())).addListener(listener);
6184
}
6285

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ccm/CCMCache.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -162,9 +162,9 @@ public void invalidate(ActionListener<Void> listener) {
162162
ClearCCMCacheAction.request(ClearCCMMessage.INSTANCE, null),
163163
ActionListener.wrap(ack -> {
164164
logger.debug("Successfully refreshed inference CCM cache for project {}.", projectResolver::getProjectId);
165-
listener.onResponse((Void) null);
165+
listener.onResponse(null);
166166
}, e -> {
167-
logger.atDebug()
167+
logger.atWarn()
168168
.withThrowable(e)
169169
.log("Failed to refresh inference CCM cache for project {}.", projectResolver::getProjectId);
170170
listener.onFailure(e);

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ccm/CCMEnablementService.java‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,8 @@ public static EnablementMetadata fromMetadata(ProjectMetadata projectMetadata) {
174174

175175
private final boolean enabled;
176176

177-
private EnablementMetadata(boolean enabled) {
177+
// Default for testing
178+
EnablementMetadata(boolean enabled) {
178179
this.enabled = enabled;
179180
}
180181

@@ -230,15 +231,11 @@ private MetadataTask(ProjectId projectId, boolean enabled, ActionListener<Acknow
230231
this.enabled = enabled;
231232
}
232233

233-
public boolean isEnabled() {
234-
return enabled;
235-
}
236-
237234
public ProjectId getProjectId() {
238235
return projectId;
239236
}
240237

241-
public EnablementMetadata executeTask(EnablementMetadata currentMetadata) {
238+
public EnablementMetadata executeTask() {
242239
return enabled ? EnablementMetadata.ENABLED : EnablementMetadata.DISABLED;
243240
}
244241
}
@@ -247,8 +244,7 @@ private static class UpdateTaskExecutor extends SimpleBatchedAckListenerTaskExec
247244
@Override
248245
public Tuple<ClusterState, ClusterStateAckListener> executeTask(MetadataTask task, ClusterState clusterState) {
249246
var projectMetadata = clusterState.metadata().getProject(task.getProjectId());
250-
var currentMetadata = EnablementMetadata.fromMetadata(projectMetadata);
251-
var updatedMetadata = task.executeTask(currentMetadata);
247+
var updatedMetadata = task.executeTask();
252248
var newProjectMetadata = ProjectMetadata.builder(projectMetadata).putCustom(EnablementMetadata.NAME, updatedMetadata);
253249
return new Tuple<>(ClusterState.builder(clusterState).putProjectMetadata(newProjectMetadata).build(), task);
254250
}

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ccm/CCMService.java‎

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -25,20 +25,22 @@ public class CCMService {
2525

2626
private final CCMPersistentStorageService ccmPersistentStorageService;
2727
private final CCMEnablementService ccmEnablementService;
28+
private final CCMCache ccmCache;
2829
private final ProjectResolver projectResolver;
2930
private final Client client;
3031

3132
public CCMService(
3233
CCMPersistentStorageService ccmPersistentStorageService,
3334
CCMEnablementService enablementService,
35+
CCMCache ccmCache,
3436
ProjectResolver projectResolver,
3537
Client client
3638
) {
3739
this.ccmPersistentStorageService = Objects.requireNonNull(ccmPersistentStorageService);
3840
this.ccmEnablementService = Objects.requireNonNull(enablementService);
41+
this.ccmCache = Objects.requireNonNull(ccmCache);
3942
this.projectResolver = Objects.requireNonNull(projectResolver);
4043
this.client = new OriginSettingClient(Objects.requireNonNull(client), ClientHelper.INFERENCE_ORIGIN);
41-
// TODO initialize the cache for the CCM configuration
4244
}
4345

4446
public void isEnabled(ActionListener<Boolean> listener) {
@@ -56,6 +58,12 @@ public void storeConfiguration(CCMModel model, ActionListener<Void> listener) {
5658
enablementListener.onFailure(e);
5759
}))
5860
)
61+
.<Void>andThen(
62+
// if we fail to invalidate the cache, it's not a big deal since the cache will eventually expire
63+
invalidateCacheListener -> ccmCache.invalidate(
64+
invalidateCacheListener.delegateResponse((delegate, e) -> delegate.onResponse(null))
65+
)
66+
)
5967
.<Void>andThen(
6068
enableAuthExecutorListener -> client.execute(
6169
AuthorizationTaskExecutor.Action.INSTANCE,
@@ -72,13 +80,10 @@ public void storeConfiguration(CCMModel model, ActionListener<Void> listener) {
7280
)
7381
)
7482
.addListener(listener);
75-
76-
// TODO invalidate the cache
7783
}
7884

7985
public void getConfiguration(ActionListener<CCMModel> listener) {
80-
// TODO get this from the cache instead
81-
ccmPersistentStorageService.get(listener);
86+
ccmCache.get(listener);
8287
}
8388

8489
public void disableCCM(ActionListener<Void> listener) {
@@ -96,6 +101,12 @@ public void disableCCM(ActionListener<Void> listener) {
96101
)
97102
)
98103
.andThen(ccmPersistentStorageService::delete)
104+
// if we fail to invalidate the cache, it's not a big deal since the cache will eventually expire
105+
.<Void>andThen(
106+
invalidateCacheListener -> ccmCache.invalidate(
107+
invalidateCacheListener.delegateResponse((delegate, e) -> delegate.onResponse(null))
108+
)
109+
)
99110
.<Void>andThen(
100111
disableAuthExecutorListener -> client.execute(
101112
AuthorizationTaskExecutor.Action.INSTANCE,
@@ -112,7 +123,5 @@ public void disableCCM(ActionListener<Void> listener) {
112123
)
113124
)
114125
.addListener(listener);
115-
116-
// TODO implement invalidating the cache
117126
}
118127
}

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/ccm/CCMAuthenticationApplierFactoryTests.java‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import org.apache.http.HttpHeaders;
1111
import org.apache.http.client.methods.HttpGet;
1212
import org.elasticsearch.ElasticsearchStatusException;
13+
import org.elasticsearch.ResourceNotFoundException;
1314
import org.elasticsearch.action.ActionListener;
1415
import org.elasticsearch.action.support.TestPlainActionFuture;
1516
import org.elasticsearch.common.settings.SecureString;
@@ -101,6 +102,64 @@ public void testGetAuthenticationApplier_ReturnsFailure_WhenConfiguringCCMIsEnab
101102
);
102103
}
103104

105+
public
106+
void
107+
testGetAuthenticationApplier_ReturnsFailure_WhenConfiguringCCMIsEnabled_TheClusterStateIndicatesItIsEnabled_ButNotConfiguration() {
108+
var ccmFeature = mock(CCMFeature.class);
109+
when(ccmFeature.isCcmSupportedEnvironment()).thenReturn(true);
110+
111+
var ccmService = mock(CCMService.class);
112+
doAnswer(invocation -> {
113+
ActionListener<Boolean> listener = invocation.getArgument(0);
114+
listener.onResponse(true);
115+
return Void.TYPE;
116+
}).when(ccmService).isEnabled(any());
117+
118+
doAnswer(invocation -> {
119+
ActionListener<CCMModel> listener = invocation.getArgument(0);
120+
listener.onFailure(new ResourceNotFoundException("not found"));
121+
return Void.TYPE;
122+
}).when(ccmService).getConfiguration(any());
123+
124+
var factory = new CCMAuthenticationApplierFactory(ccmFeature, ccmService);
125+
var listener = new TestPlainActionFuture<CCMAuthenticationApplierFactory.AuthApplier>();
126+
factory.getAuthenticationApplier(listener);
127+
128+
var exception = expectThrows(ElasticsearchStatusException.class, () -> listener.actionGet(TimeValue.THIRTY_SECONDS));
129+
assertThat(
130+
exception.getMessage(),
131+
containsString(
132+
"Cloud connected mode configuration is in an inconsistent state. "
133+
+ "Please try configuring it again using PUT _inference/_ccm"
134+
)
135+
);
136+
}
137+
138+
public void testGetAuthenticationApplier_ReturnsFailure_ConfiguringCCMIsEnabled_ClusterStateEnabled_FailureGettingConfiguration() {
139+
var ccmFeature = mock(CCMFeature.class);
140+
when(ccmFeature.isCcmSupportedEnvironment()).thenReturn(true);
141+
142+
var ccmService = mock(CCMService.class);
143+
doAnswer(invocation -> {
144+
ActionListener<Boolean> listener = invocation.getArgument(0);
145+
listener.onResponse(true);
146+
return Void.TYPE;
147+
}).when(ccmService).isEnabled(any());
148+
149+
doAnswer(invocation -> {
150+
ActionListener<CCMModel> listener = invocation.getArgument(0);
151+
listener.onFailure(new IllegalArgumentException("not found"));
152+
return Void.TYPE;
153+
}).when(ccmService).getConfiguration(any());
154+
155+
var factory = new CCMAuthenticationApplierFactory(ccmFeature, ccmService);
156+
var listener = new TestPlainActionFuture<CCMAuthenticationApplierFactory.AuthApplier>();
157+
factory.getAuthenticationApplier(listener);
158+
159+
var exception = expectThrows(IllegalArgumentException.class, () -> listener.actionGet(TimeValue.THIRTY_SECONDS));
160+
assertThat(exception.getMessage(), containsString("not found"));
161+
}
162+
104163
public void testGetAuthenticationApplier_ReturnsApiKey_WhenConfiguringCCMIsEnabled_AndSet() {
105164
var secret = "secret";
106165

0 commit comments

Comments
 (0)