Add tests for MLAgent Get and Delete (opensearch-project#1791)

Signed-off-by: Mingshi Liu <[email protected]>
Zhangxunmt · Dec 20, 2023 · 9ba7c60 · 9ba7c60
1 parent 9aea9a0
commit 9ba7c60
Show file tree

Hide file tree

Showing 3 changed files with 112 additions and 3 deletions.
diff --git a/common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentDeleteRequestTest.java b/common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentDeleteRequestTest.java
@@ -5,10 +5,13 @@
 package org.opensearch.ml.common.transport.agent;
 
 import org.junit.Test;
+import org.opensearch.action.ActionRequest;
 import org.opensearch.action.ActionRequestValidationException;
 import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.core.common.io.stream.StreamOutput;
 
 import java.io.IOException;
+import java.io.UncheckedIOException;
 
 import static org.junit.Assert.assertEquals;
 import static org.opensearch.action.ValidateActions.addValidationError;
@@ -56,10 +59,46 @@ public void validate_Failure() {
     }
 
     @Test
-    public void fromActionRequest() throws IOException {
+    public void fromActionRequest_Success() throws IOException {
         agentId = "test-lmn";
         MLAgentDeleteRequest mLAgentDeleteRequest = new MLAgentDeleteRequest(agentId);
         assertEquals(mLAgentDeleteRequest.fromActionRequest(mLAgentDeleteRequest), mLAgentDeleteRequest);
 
     }
+    @Test
+    public void fromActionRequest_Success_fromActionRequest() throws IOException {
+        agentId = "test-opq";
+        MLAgentDeleteRequest mLAgentDeleteRequest = new MLAgentDeleteRequest(agentId);
+
+        ActionRequest actionRequest = new ActionRequest() {
+            @Override
+            public ActionRequestValidationException validate() {
+                return null;
+            }
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {
+                mLAgentDeleteRequest.writeTo(out);
+            }
+        };
+        MLAgentDeleteRequest request = mLAgentDeleteRequest.fromActionRequest(actionRequest);
+        assertEquals(request.agentId, agentId);
+    }
+
+    @Test(expected = UncheckedIOException.class)
+    public void fromActionRequest_IOException() {
+        agentId = "test-rst";
+        MLAgentDeleteRequest mLAgentDeleteRequest = new MLAgentDeleteRequest(agentId);
+        ActionRequest actionRequest = new ActionRequest() {
+            @Override
+            public ActionRequestValidationException validate() {
+                return null;
+            }
+
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {
+                throw new IOException();
+            }
+        };
+        mLAgentDeleteRequest.fromActionRequest(actionRequest);
+    }
 }
diff --git a/common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentGetRequestTest.java b/common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentGetRequestTest.java
@@ -5,10 +5,14 @@
 package org.opensearch.ml.common.transport.agent;
 
 import org.junit.Test;
+import org.opensearch.action.ActionRequest;
 import org.opensearch.action.ActionRequestValidationException;
 import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.core.common.io.stream.StreamOutput;
 
 import java.io.IOException;
+import java.io.UncheckedIOException;
+
 import static org.junit.Assert.assertEquals;
 import static org.opensearch.action.ValidateActions.addValidationError;
 
@@ -54,11 +58,48 @@ public void validate_Failure() {
         mLAgentGetRequest.validate().equals(exception) ;
     }
     @Test
-    public void fromActionRequest() throws IOException {
+    public void fromActionRequest_Success()  throws IOException {
         agentId = "test-lmn";
         MLAgentGetRequest mLAgentGetRequest = new MLAgentGetRequest(agentId);
         assertEquals(mLAgentGetRequest.fromActionRequest(mLAgentGetRequest), mLAgentGetRequest);
     }
+
+    @Test
+    public void fromActionRequest_Success_fromActionRequest() throws IOException {
+        agentId = "test-opq";
+        MLAgentGetRequest mLAgentGetRequest = new MLAgentGetRequest(agentId);
+
+        ActionRequest actionRequest = new ActionRequest() {
+            @Override
+            public ActionRequestValidationException validate() {
+                return null;
+            }
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {
+                mLAgentGetRequest.writeTo(out);
+            }
+        };
+        MLAgentGetRequest request = mLAgentGetRequest.fromActionRequest(actionRequest);
+        assertEquals(request.agentId, agentId);
+    }
+
+    @Test(expected = UncheckedIOException.class)
+    public void fromActionRequest_IOException() {
+        agentId = "test-rst";
+        MLAgentGetRequest mLAgentGetRequest = new MLAgentGetRequest(agentId);
+        ActionRequest actionRequest = new ActionRequest() {
+            @Override
+            public ActionRequestValidationException validate() {
+                return null;
+            }
+
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {
+                throw new IOException();
+            }
+        };
+        mLAgentGetRequest.fromActionRequest(actionRequest);
+    }
 }
 
 
diff --git a/common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentGetResponseTest.java b/common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentGetResponseTest.java
@@ -99,11 +99,40 @@ public void toXContent() throws IOException {
     }
 
     @Test
-    public void FromActionResponse() throws IOException {
+    public void fromActionResponse_Success() throws IOException {
         MLAgentGetResponse mlAgentGetResponse = MLAgentGetResponse.builder()
                 .mlAgent(mlAgent)
                 .build();
         assertEquals(mlAgentGetResponse.fromActionResponse(mlAgentGetResponse), mlAgentGetResponse);
 
         }
+    @Test
+    public void fromActionResponse_Success_fromActionResponse() throws IOException {
+        MLAgentGetResponse mlAgentGetResponse = MLAgentGetResponse.builder()
+                .mlAgent(mlAgent)
+                .build();
+
+        ActionResponse actionResponse = new ActionResponse() {
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {
+                mlAgentGetResponse.writeTo(out);
+            }
+        };
+        MLAgentGetResponse response = mlAgentGetResponse.fromActionResponse(actionResponse);
+        assertEquals(response.getMlAgent(), mlAgent);
+    }
+
+    @Test(expected = UncheckedIOException.class)
+    public void fromActionResponse_IOException() {
+        MLAgentGetResponse mlAgentGetResponse = MLAgentGetResponse.builder()
+                .mlAgent(mlAgent)
+                .build();
+        ActionResponse actionResponse = new ActionResponse() {
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {
+                throw new IOException();
+            }
+        };
+        mlAgentGetResponse.fromActionResponse(actionResponse);
+    }
     }