Skip to content

Commit 2e3246c

Browse files
jonathan-buttnerjonathan-buttner
committed
Starting tests
1 parent a7f0f91 commit 2e3246c

File tree

5 files changed

+252
-41
lines changed

5 files changed

+252
-41
lines changed

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

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353
import org.elasticsearch.xpack.inference.LocalStateInferencePlugin;
5454
import org.elasticsearch.xpack.inference.model.TestModel;
5555
import org.elasticsearch.xpack.inference.registry.ModelRegistry;
56+
import org.elasticsearch.xpack.inference.registry.ModelStoreResponse;
5657
import org.elasticsearch.xpack.inference.services.elasticsearch.ElasticsearchInternalModel;
5758
import org.elasticsearch.xpack.inference.services.elasticsearch.ElasticsearchInternalService;
5859
import org.elasticsearch.xpack.inference.services.elasticsearch.ElserInternalServiceSettingsTests;
@@ -600,7 +601,7 @@ public void testGetUnparsedModelMap_ThrowsResourceNotFound_WhenNoHitsReturned()
600601
}
601602

602603
public void testStoreModels_ReturnsEmptyList_WhenGivenNoModelsToStore() {
603-
PlainActionFuture<List<ModelRegistry.ModelStoreResponse>> storeListener = new PlainActionFuture<>();
604+
PlainActionFuture<List<ModelStoreResponse>> storeListener = new PlainActionFuture<>();
604605
modelRegistry.storeModels(List.of(), storeListener, TimeValue.THIRTY_SECONDS);
605606

606607
var response = storeListener.actionGet(TimeValue.THIRTY_SECONDS);
@@ -620,12 +621,12 @@ public void testStoreModels_StoresSingleInferenceEndpoint() {
620621
new TestModel.TestSecretSettings(secrets)
621622
);
622623

623-
PlainActionFuture<List<ModelRegistry.ModelStoreResponse>> storeListener = new PlainActionFuture<>();
624+
PlainActionFuture<List<ModelStoreResponse>> storeListener = new PlainActionFuture<>();
624625
modelRegistry.storeModels(List.of(model), storeListener, TimeValue.THIRTY_SECONDS);
625626

626627
var response = storeListener.actionGet(TimeValue.THIRTY_SECONDS);
627628
assertThat(response.size(), Matchers.is(1));
628-
assertThat(response.get(0), Matchers.is(new ModelRegistry.ModelStoreResponse(inferenceId, RestStatus.CREATED, null)));
629+
assertThat(response.get(0), Matchers.is(new ModelStoreResponse(inferenceId, RestStatus.CREATED, null)));
629630

630631
assertMinimalServiceSettings(modelRegistry, model);
631632

@@ -659,13 +660,13 @@ public void testStoreModels_StoresMultipleInferenceEndpoints() {
659660
new TestModel.TestSecretSettings(secrets)
660661
);
661662

662-
PlainActionFuture<List<ModelRegistry.ModelStoreResponse>> storeListener = new PlainActionFuture<>();
663+
PlainActionFuture<List<ModelStoreResponse>> storeListener = new PlainActionFuture<>();
663664
modelRegistry.storeModels(List.of(model1, model2), storeListener, TimeValue.THIRTY_SECONDS);
664665

665666
var response = storeListener.actionGet(TimeValue.THIRTY_SECONDS);
666667
assertThat(response.size(), Matchers.is(2));
667-
assertThat(response.get(0), Matchers.is(new ModelRegistry.ModelStoreResponse(inferenceId1, RestStatus.CREATED, null)));
668-
assertThat(response.get(1), Matchers.is(new ModelRegistry.ModelStoreResponse(inferenceId2, RestStatus.CREATED, null)));
668+
assertThat(response.get(0), Matchers.is(new ModelStoreResponse(inferenceId1, RestStatus.CREATED, null)));
669+
assertThat(response.get(1), Matchers.is(new ModelStoreResponse(inferenceId2, RestStatus.CREATED, null)));
669670

670671
assertModelAndMinimalSettingsWithSecrets(modelRegistry, model1, secrets);
671672
assertModelAndMinimalSettingsWithSecrets(modelRegistry, model2, secrets);
@@ -716,12 +717,12 @@ public void testStoreModels_StoresOneModel_FailsToStoreSecond_WhenVersionConflic
716717
new TestModel.TestSecretSettings(secrets)
717718
);
718719

719-
PlainActionFuture<List<ModelRegistry.ModelStoreResponse>> storeListener = new PlainActionFuture<>();
720+
PlainActionFuture<List<ModelStoreResponse>> storeListener = new PlainActionFuture<>();
720721
modelRegistry.storeModels(List.of(model1, model2), storeListener, TimeValue.THIRTY_SECONDS);
721722

722723
var response = storeListener.actionGet(TimeValue.THIRTY_SECONDS);
723724
assertThat(response.size(), Matchers.is(2));
724-
assertThat(response.get(0), Matchers.is(new ModelRegistry.ModelStoreResponse(inferenceId, RestStatus.CREATED, null)));
725+
assertThat(response.get(0), Matchers.is(new ModelStoreResponse(inferenceId, RestStatus.CREATED, null)));
725726
assertThat(response.get(1).inferenceId(), Matchers.is(model2.getInferenceEntityId()));
726727
assertThat(response.get(1).status(), Matchers.is(RestStatus.CONFLICT));
727728
assertTrue(response.get(1).failed());
@@ -758,12 +759,12 @@ public void testStoreModels_StoresOneModel_RemovesSecondDuplicateModelFromList_D
758759
new TestModel.TestSecretSettings(secrets)
759760
);
760761

761-
PlainActionFuture<List<ModelRegistry.ModelStoreResponse>> storeListener = new PlainActionFuture<>();
762+
PlainActionFuture<List<ModelStoreResponse>> storeListener = new PlainActionFuture<>();
762763
modelRegistry.storeModels(List.of(model1, model1, model2), storeListener, TimeValue.THIRTY_SECONDS);
763764

764765
var response = storeListener.actionGet(TimeValue.THIRTY_SECONDS);
765766
assertThat(response.size(), Matchers.is(1));
766-
assertThat(response.get(0), Matchers.is(new ModelRegistry.ModelStoreResponse(inferenceId, RestStatus.CREATED, null)));
767+
assertThat(response.get(0), Matchers.is(new ModelStoreResponse(inferenceId, RestStatus.CREATED, null)));
767768

768769
assertModelAndMinimalSettingsWithSecrets(modelRegistry, model1, secrets);
769770
assertIndicesContainExpectedDocsCount(model1, 2);
@@ -783,7 +784,7 @@ public void testStoreModels_FailsToStoreModel_WhenInferenceIndexDocumentAlreadyE
783784

784785
storeCorruptedModel(model, false);
785786

786-
PlainActionFuture<List<ModelRegistry.ModelStoreResponse>> storeListener = new PlainActionFuture<>();
787+
PlainActionFuture<List<ModelStoreResponse>> storeListener = new PlainActionFuture<>();
787788
modelRegistry.storeModels(List.of(model), storeListener, TimeValue.THIRTY_SECONDS);
788789

789790
var response = storeListener.actionGet(TimeValue.THIRTY_SECONDS);
@@ -837,7 +838,7 @@ public void testStoreModels_OnFailure_RemovesPartialWritesOfInferenceEndpoint()
837838
storeCorruptedModel(model1, false);
838839
storeCorruptedModel(model2, true);
839840

840-
PlainActionFuture<List<ModelRegistry.ModelStoreResponse>> storeListener = new PlainActionFuture<>();
841+
PlainActionFuture<List<ModelStoreResponse>> storeListener = new PlainActionFuture<>();
841842
modelRegistry.storeModels(List.of(model1, model2, model3), storeListener, TimeValue.THIRTY_SECONDS);
842843

843844
var response = storeListener.actionGet(TimeValue.THIRTY_SECONDS);

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

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -253,11 +253,6 @@ private static Map<String, DefaultModelConfig> initDefaultEndpoints(
253253
);
254254
}
255255

256-
@Override
257-
public void onNodeStarted() {
258-
// authorizationHandler.init();
259-
}
260-
261256
@Override
262257
protected void validateRerankParameters(Boolean returnDocuments, Integer topN, ValidationException validationException) {
263258
if (returnDocuments != null) {

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

Lines changed: 10 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,6 @@
3232
import java.util.Map;
3333
import java.util.Objects;
3434
import java.util.Set;
35-
import java.util.concurrent.CountDownLatch;
36-
import java.util.concurrent.TimeUnit;
3735
import java.util.concurrent.atomic.AtomicBoolean;
3836
import java.util.concurrent.atomic.AtomicReference;
3937
import java.util.stream.Collectors;
@@ -50,7 +48,6 @@ public class AuthorizationPoller extends AllocatedPersistentTask {
5048
private final ServiceComponents serviceComponents;
5149
private final ModelRegistry modelRegistry;
5250
private final ElasticInferenceServiceAuthorizationRequestHandler authorizationHandler;
53-
private final CountDownLatch firstAuthorizationCompletedLatch = new CountDownLatch(1);
5451
private final Sender sender;
5552
private final Runnable callback;
5653
private final AtomicReference<Scheduler.ScheduledCancellable> lastAuthTask = new AtomicReference<>(null);
@@ -117,23 +114,13 @@ public void start() {
117114
}
118115
}
119116

120-
/**
121-
* Waits the specified amount of time for the first authorization call to complete. This is mainly to make testing easier.
122-
* @param waitTime the max time to wait
123-
* @throws IllegalStateException if the wait time is exceeded or the call receives an {@link InterruptedException}
124-
*/
125-
public void waitForAuthorizationToComplete(TimeValue waitTime) {
126-
try {
127-
if (firstAuthorizationCompletedLatch.await(waitTime.getSeconds(), TimeUnit.SECONDS) == false) {
128-
throw new IllegalStateException("The wait time has expired for authorization to complete.");
129-
}
130-
} catch (InterruptedException e) {
131-
throw new IllegalStateException("Waiting for authorization to complete was interrupted");
132-
}
133-
}
134-
135117
@Override
136118
protected void onCancelled() {
119+
shutdown();
120+
}
121+
122+
// default for testing
123+
void shutdown() {
137124
shutdown.set(true);
138125
if (lastAuthTask.get() != null) {
139126
lastAuthTask.get().cancel();
@@ -182,7 +169,8 @@ private void scheduleAndSendAuthorizationRequest() {
182169
sendAuthorizationRequest();
183170
}
184171

185-
private void sendAuthorizationRequest() {
172+
// default for testing
173+
void sendAuthorizationRequest() {
186174
if (modelRegistry.isReady() == false) {
187175
return;
188176
}
@@ -191,7 +179,6 @@ private void sendAuthorizationRequest() {
191179
if (callback != null) {
192180
callback.run();
193181
}
194-
firstAuthorizationCompletedLatch.countDown();
195182
}).delegateResponse((delegate, e) -> {
196183
logger.atWarn().withThrowable(e).log("Failed processing EIS preconfigured endpoints");
197184
delegate.onResponse(null);
@@ -227,11 +214,11 @@ private void storePreconfiguredModels(Set<String> newInferenceIds, ActionListene
227214
return;
228215
}
229216

230-
logger.debug("Storing new EIS preconfigured inference endpoints with inference IDs {}", newInferenceIds);
217+
logger.info("Storing new EIS preconfigured inference endpoints with inference IDs {}", newInferenceIds);
231218
var modelsToAdd = PreconfiguredEndpointModelAdapter.getModels(newInferenceIds, elasticInferenceServiceComponents);
232219
var storeRequest = new StoreInferenceEndpointsAction.Request(modelsToAdd, TimeValue.THIRTY_SECONDS);
233220

234-
ActionListener<StoreInferenceEndpointsAction.Response> storeListener = ActionListener.wrap(responses -> {
221+
ActionListener<StoreInferenceEndpointsAction.Response> logResultsListener = ActionListener.wrap(responses -> {
235222
for (var response : responses.getResults()) {
236223
if (response.failed()) {
237224
logger.atWarn()
@@ -247,7 +234,7 @@ private void storePreconfiguredModels(Set<String> newInferenceIds, ActionListene
247234
client.execute(
248235
StoreInferenceEndpointsAction.INSTANCE,
249236
storeRequest,
250-
ActionListener.runAfter(storeListener, () -> listener.onResponse(null))
237+
ActionListener.runAfter(logResultsListener, () -> listener.onResponse(null))
251238
);
252239
}
253240
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ public static List<Model> getModels(Set<String> inferenceIds, ElasticInferenceSe
2727
.toList();
2828
}
2929

30-
private static Model createModel(
30+
public static Model createModel(
3131
InternalPreconfiguredEndpoints.MinimalModel minimalModel,
3232
ElasticInferenceServiceComponents elasticInferenceServiceComponents
3333
) {

0 commit comments

Comments
 (0)