|
@@ -12,19 +12,16 @@ import org.elasticsearch.common.bytes.BytesArray;
|
|
|
import org.elasticsearch.rest.RestRequest;
|
|
|
import org.elasticsearch.test.rest.FakeRestRequest;
|
|
|
import org.elasticsearch.test.rest.RestActionTestCase;
|
|
|
-import org.elasticsearch.xcontent.XContentBuilder;
|
|
|
import org.elasticsearch.xcontent.XContentType;
|
|
|
import org.elasticsearch.xpack.core.ml.action.CreateTrainedModelAssignmentAction;
|
|
|
import org.elasticsearch.xpack.core.ml.action.UpdateTrainedModelDeploymentAction;
|
|
|
|
|
|
-import java.io.IOException;
|
|
|
import java.util.HashMap;
|
|
|
|
|
|
+import static org.elasticsearch.rest.RestResponseUtils.setUpXContentMock;
|
|
|
import static org.hamcrest.Matchers.equalTo;
|
|
|
import static org.hamcrest.Matchers.instanceOf;
|
|
|
-import static org.mockito.ArgumentMatchers.any;
|
|
|
import static org.mockito.Mockito.mock;
|
|
|
-import static org.mockito.Mockito.when;
|
|
|
|
|
|
public class RestUpdateTrainedModelDeploymentActionTests extends RestActionTestCase {
|
|
|
public void testNumberOfAllocationInParam() {
|
|
@@ -37,7 +34,7 @@ public class RestUpdateTrainedModelDeploymentActionTests extends RestActionTestC
|
|
|
assertEquals(request.getNumberOfAllocations().intValue(), 5);
|
|
|
|
|
|
executeCalled.set(true);
|
|
|
- return newMockResponse();
|
|
|
+ return setUpXContentMock(mock(CreateTrainedModelAssignmentAction.Response.class));
|
|
|
}));
|
|
|
var params = new HashMap<String, String>();
|
|
|
params.put("number_of_allocations", "5");
|
|
@@ -60,7 +57,7 @@ public class RestUpdateTrainedModelDeploymentActionTests extends RestActionTestC
|
|
|
assertEquals(request.getNumberOfAllocations().intValue(), 6);
|
|
|
|
|
|
executeCalled.set(true);
|
|
|
- return newMockResponse();
|
|
|
+ return setUpXContentMock(mock(CreateTrainedModelAssignmentAction.Response.class));
|
|
|
}));
|
|
|
|
|
|
final String content = """
|
|
@@ -73,16 +70,4 @@ public class RestUpdateTrainedModelDeploymentActionTests extends RestActionTestC
|
|
|
dispatchRequest(inferenceRequest);
|
|
|
assertThat(executeCalled.get(), equalTo(true));
|
|
|
}
|
|
|
-
|
|
|
- private static CreateTrainedModelAssignmentAction.Response newMockResponse() {
|
|
|
- final var response = mock(CreateTrainedModelAssignmentAction.Response.class);
|
|
|
- try {
|
|
|
- when(response.toXContent(any(), any())).thenAnswer(
|
|
|
- invocation -> asInstanceOf(XContentBuilder.class, invocation.getArgument(0)).startObject().endObject()
|
|
|
- );
|
|
|
- } catch (IOException e) {
|
|
|
- fail(e);
|
|
|
- }
|
|
|
- return response;
|
|
|
- }
|
|
|
}
|