Skip to content

Commit e16cb95

Browse files
dbwiddisdbwiddis
committed
Add integration test covering bulk API
Signed-off-by: Daniel Widdis <[email protected]>
1 parent 7eabdd5 commit e16cb95

File tree

2 files changed

+154
-0
lines changed

2 files changed

+154
-0
lines changed

‎plugin/src/test/java/org/opensearch/ml/rest/MLCommonsTenantAwareRestTestCase.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,8 @@ public abstract class MLCommonsTenantAwareRestTestCase extends MLCommonsRestTest
6262
// REST Response error reasons
6363
protected static final String MISSING_TENANT_REASON = "Tenant ID header is missing";
6464
protected static final String NO_PERMISSION_REASON = "You don't have permission to access this resource";
65+
protected static final String DEPLOYED_REASON =
66+
"Model cannot be deleted in deploying or deployed state. Try undeploy model first then delete";
6567

6668
// Common constants and fields used in subclasses
6769
protected static final String CONNECTOR_ID = "connector_id";
@@ -167,6 +169,10 @@ protected static SearchResponse searchResponseFromResponse(Response response) th
167169
return SearchResponse.fromXContent(parser);
168170
}
169171

172+
protected static void assertBadRequest(Response response) {
173+
assertEquals(RestStatus.BAD_REQUEST.getStatus(), response.getStatusLine().getStatusCode());
174+
}
175+
170176
protected static void assertNotFound(Response response) {
171177
assertEquals(RestStatus.NOT_FOUND.getStatus(), response.getStatusLine().getStatusCode());
172178
}

‎plugin/src/test/java/org/opensearch/ml/rest/RestMLModelUndeployTenantAwareIT.java

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.rest;
7+
8+
import static org.opensearch.ml.common.CommonValue.TENANT_ID;
9+
import static org.opensearch.ml.common.MLModel.MODEL_STATE_FIELD;
10+
import static org.opensearch.ml.common.MLTask.MODEL_ID_FIELD;
11+
import static org.opensearch.ml.rest.RestMLRAGSearchProcessorIT.COHERE_CONNECTOR_BLUEPRINT;
12+
13+
import java.util.Map;
14+
import java.util.concurrent.TimeUnit;
15+
16+
import org.opensearch.client.Response;
17+
import org.opensearch.client.ResponseException;
18+
import org.opensearch.rest.RestRequest;
19+
20+
public class RestMLModelUndeployTenantAwareIT extends MLCommonsTenantAwareRestTestCase {
21+
22+
// Tests the client.bulk API used for undeploying models
23+
public void testModelDeployUndeploy() throws Exception {
24+
boolean multiTenancyEnabled = isMultiTenancyEnabled();
25+
26+
/*
27+
* Setup
28+
*/
29+
// Create a connector to use
30+
RestRequest createConnectorRequest = getRestRequestWithHeadersAndContent(tenantId, COHERE_CONNECTOR_BLUEPRINT);
31+
Response response = makeRequest(createConnectorRequest, POST, CONNECTORS_PATH + "_create");
32+
assertOK(response);
33+
Map<String, Object> map = responseToMap(response);
34+
assertTrue(map.containsKey(CONNECTOR_ID));
35+
String connectorId = map.get(CONNECTOR_ID).toString();
36+
37+
/*
38+
* Create
39+
*/
40+
// Register and deploy a remote model with a tenant id
41+
RestRequest registerModelRequest = getRestRequestWithHeadersAndContent(
42+
tenantId,
43+
registerRemoteModelContent("test model", connectorId, null)
44+
);
45+
response = makeRequest(registerModelRequest, POST, MODELS_PATH + "_register?deploy=true");
46+
assertOK(response);
47+
map = responseToMap(response);
48+
assertTrue(map.containsKey(MODEL_ID_FIELD));
49+
String modelId = map.get(MODEL_ID_FIELD).toString();
50+
51+
/*
52+
* Get
53+
*/
54+
// Now get that model and confirm it's deployed
55+
assertBusy(() -> {
56+
Response getResponse = makeRequest(tenantRequest, GET, MODELS_PATH + modelId);
57+
assertOK(getResponse);
58+
Map<String, Object> responseMap = responseToMap(getResponse);
59+
assertEquals("DEPLOYED", responseMap.get(MODEL_STATE_FIELD).toString());
60+
if (multiTenancyEnabled) {
61+
assertEquals(tenantId, responseMap.get(TENANT_ID));
62+
} else {
63+
assertNull(responseMap.get(TENANT_ID));
64+
}
65+
}, 20, TimeUnit.SECONDS);
66+
67+
/*
68+
* Test delete/deploy interaction
69+
*/
70+
// Attempt to delete, should fail because it's deployed
71+
ResponseException ex = assertThrows(ResponseException.class, () -> makeRequest(tenantRequest, DELETE, MODELS_PATH + modelId));
72+
response = ex.getResponse();
73+
assertBadRequest(response);
74+
map = responseToMap(response);
75+
assertEquals(DEPLOYED_REASON, getErrorReasonFromResponseMap(map));
76+
77+
// Verify still exists
78+
response = makeRequest(tenantRequest, GET, MODELS_PATH + modelId);
79+
assertOK(response);
80+
81+
/*
82+
* Undeploy
83+
*/
84+
// Undeploy the model which uses the bulk API
85+
if (multiTenancyEnabled) {
86+
// Try with the wrong tenant
87+
ex = assertThrows(ResponseException.class, () -> makeRequest(otherTenantRequest, POST, MODELS_PATH + modelId + "/_undeploy"));
88+
response = ex.getResponse();
89+
map = responseToMap(response);
90+
if (DDB) {
91+
assertNotFound(response);
92+
assertTrue(getErrorReasonFromResponseMap(map).startsWith("Failed to find model"));
93+
} else {
94+
assertForbidden(response);
95+
assertEquals(NO_PERMISSION_REASON, getErrorReasonFromResponseMap(map));
96+
}
97+
98+
// Try with a null tenant
99+
ex = assertThrows(ResponseException.class, () -> makeRequest(nullTenantRequest, POST, MODELS_PATH + modelId + "/_undeploy"));
100+
response = ex.getResponse();
101+
assertForbidden(response);
102+
map = responseToMap(response);
103+
assertEquals(MISSING_TENANT_REASON, getErrorReasonFromResponseMap(map));
104+
}
105+
106+
// Now do with correct tenant
107+
response = makeRequest(tenantRequest, POST, MODELS_PATH + modelId + "/_undeploy");
108+
assertOK(response);
109+
// This is an MLUndeployControllerNodeResponse
110+
map = responseToMap(response);
111+
// This map's keys are the nodes, and the values are a map with "stats" key
112+
// One of these is a map object with modelId as key and "undeployed" as value
113+
String expectedValue = modelId + "=undeployed";
114+
assertTrue(map.toString().contains(expectedValue));
115+
116+
// Verify the undeploy update
117+
response = makeRequest(tenantRequest, GET, MODELS_PATH + modelId);
118+
assertOK(response);
119+
map = responseToMap(response);
120+
assertEquals("UNDEPLOYED", map.get(MODEL_STATE_FIELD).toString());
121+
if (multiTenancyEnabled) {
122+
assertEquals(tenantId, map.get(TENANT_ID));
123+
} else {
124+
assertNull(map.get(TENANT_ID));
125+
}
126+
127+
/*
128+
* Delete
129+
*/
130+
// Delete, should now succeed because it's deployed
131+
response = makeRequest(tenantRequest, DELETE, MODELS_PATH + modelId);
132+
assertOK(response);
133+
map = responseToMap(response);
134+
assertEquals(modelId, map.get(DOC_ID).toString());
135+
136+
// Verify the deletion
137+
ex = assertThrows(ResponseException.class, () -> makeRequest(tenantRequest, GET, MODELS_PATH + modelId));
138+
response = ex.getResponse();
139+
assertNotFound(response);
140+
map = responseToMap(response);
141+
assertEquals("Failed to find model with the provided model id: " + modelId, getErrorReasonFromResponseMap(map));
142+
143+
/*
144+
* Cleanup other resources created
145+
*/
146+
deleteAndWaitForSearch(tenantId, CONNECTORS_PATH, connectorId, 0);
147+
}
148+
}

0 commit comments

Comments
 (0)