From ddf5b5cd76301fc1a18cb92616da6ae12c3cdd4e Mon Sep 17 00:00:00 2001 From: Ryan Bogan Date: Tue, 2 Jan 2024 11:27:52 -0800 Subject: [PATCH] Localize static mocked variable Signed-off-by: Ryan Bogan --- .../transport/GetModelResponseTests.java | 55 +++++++++---------- 1 file changed, 27 insertions(+), 28 deletions(-) diff --git a/src/test/java/org/opensearch/knn/plugin/transport/GetModelResponseTests.java b/src/test/java/org/opensearch/knn/plugin/transport/GetModelResponseTests.java index a9a1a0d8f..04c94de7e 100644 --- a/src/test/java/org/opensearch/knn/plugin/transport/GetModelResponseTests.java +++ b/src/test/java/org/opensearch/knn/plugin/transport/GetModelResponseTests.java @@ -11,7 +11,6 @@ package org.opensearch.knn.plugin.transport; -import org.junit.BeforeClass; import org.mockito.MockedStatic; import org.opensearch.Version; import org.opensearch.common.io.stream.BytesStreamOutput; @@ -33,16 +32,6 @@ public class GetModelResponseTests extends KNNTestCase { - private static MockedStatic knnClusterUtilMockedStatic; - - @BeforeClass - public static void setup() { - final KNNClusterUtil knnClusterUtil = mock(KNNClusterUtil.class); - knnClusterUtilMockedStatic = mockStatic(KNNClusterUtil.class); - when(knnClusterUtil.getClusterMinVersion()).thenReturn(Version.CURRENT); - knnClusterUtilMockedStatic.when(KNNClusterUtil::instance).thenReturn(knnClusterUtil); - } - private ModelMetadata getModelMetadata(ModelState state) { return new ModelMetadata(KNNEngine.DEFAULT, SpaceType.DEFAULT, 4, state, "2021-03-27 10:15:30 AM +05:30", "test model", "", ""); } @@ -59,25 +48,35 @@ public void testStreams() throws IOException { } public void testXContent() throws IOException { - String modelId = "test-model"; - byte[] testModelBlob = "hello".getBytes(); - Model model = new Model(getModelMetadata(ModelState.CREATED), testModelBlob, modelId); - GetModelResponse getModelResponse = new GetModelResponse(model); - String expectedResponseString = - "{\"model_id\":\"test-model\",\"model_blob\":\"aGVsbG8=\",\"state\":\"created\",\"timestamp\":\"2021-03-27 10:15:30 AM +05:30\",\"description\":\"test model\",\"error\":\"\",\"space_type\":\"l2\",\"dimension\":4,\"engine\":\"nmslib\",\"training_node_assignment\":\"\"}"; - XContentBuilder xContentBuilder = XContentFactory.jsonBuilder(); - getModelResponse.toXContent(xContentBuilder, null); - assertEquals(expectedResponseString, xContentBuilder.toString()); + try (MockedStatic knnClusterUtilMockedStatic = mockStatic(KNNClusterUtil.class)) { + final KNNClusterUtil knnClusterUtil = mock(KNNClusterUtil.class); + when(knnClusterUtil.getClusterMinVersion()).thenReturn(Version.CURRENT); + knnClusterUtilMockedStatic.when(KNNClusterUtil::instance).thenReturn(knnClusterUtil); + String modelId = "test-model"; + byte[] testModelBlob = "hello".getBytes(); + Model model = new Model(getModelMetadata(ModelState.CREATED), testModelBlob, modelId); + GetModelResponse getModelResponse = new GetModelResponse(model); + String expectedResponseString = + "{\"model_id\":\"test-model\",\"model_blob\":\"aGVsbG8=\",\"state\":\"created\",\"timestamp\":\"2021-03-27 10:15:30 AM +05:30\",\"description\":\"test model\",\"error\":\"\",\"space_type\":\"l2\",\"dimension\":4,\"engine\":\"nmslib\",\"training_node_assignment\":\"\"}"; + XContentBuilder xContentBuilder = XContentFactory.jsonBuilder(); + getModelResponse.toXContent(xContentBuilder, null); + assertEquals(expectedResponseString, xContentBuilder.toString()); + } } public void testXContentWithNoModelBlob() throws IOException { - String modelId = "test-model"; - Model model = new Model(getModelMetadata(ModelState.FAILED), null, modelId); - GetModelResponse getModelResponse = new GetModelResponse(model); - String expectedResponseString = - "{\"model_id\":\"test-model\",\"model_blob\":\"\",\"state\":\"failed\",\"timestamp\":\"2021-03-27 10:15:30 AM +05:30\",\"description\":\"test model\",\"error\":\"\",\"space_type\":\"l2\",\"dimension\":4,\"engine\":\"nmslib\",\"training_node_assignment\":\"\"}"; - XContentBuilder xContentBuilder = XContentFactory.jsonBuilder(); - getModelResponse.toXContent(xContentBuilder, null); - assertEquals(expectedResponseString, xContentBuilder.toString()); + try (MockedStatic knnClusterUtilMockedStatic = mockStatic(KNNClusterUtil.class)) { + final KNNClusterUtil knnClusterUtil = mock(KNNClusterUtil.class); + when(knnClusterUtil.getClusterMinVersion()).thenReturn(Version.CURRENT); + knnClusterUtilMockedStatic.when(KNNClusterUtil::instance).thenReturn(knnClusterUtil); + String modelId = "test-model"; + Model model = new Model(getModelMetadata(ModelState.FAILED), null, modelId); + GetModelResponse getModelResponse = new GetModelResponse(model); + String expectedResponseString = + "{\"model_id\":\"test-model\",\"model_blob\":\"\",\"state\":\"failed\",\"timestamp\":\"2021-03-27 10:15:30 AM +05:30\",\"description\":\"test model\",\"error\":\"\",\"space_type\":\"l2\",\"dimension\":4,\"engine\":\"nmslib\",\"training_node_assignment\":\"\"}"; + XContentBuilder xContentBuilder = XContentFactory.jsonBuilder(); + getModelResponse.toXContent(xContentBuilder, null); + assertEquals(expectedResponseString, xContentBuilder.toString()); + } } }