From a3371fe40d00e78d9ee6a134383b4aca6ac74adf Mon Sep 17 00:00:00 2001
From: Bhavana Goud Ramaram <rbhavna@amazon.com>
Date: Fri, 10 Jan 2025 19:43:29 -0600
Subject: [PATCH] add back auto redeploy ITs from reverted commits

Signed-off-by: Bhavana Goud Ramaram <rbhavna@amazon.com>
---
 .../autoredeploy/MLModelAutoReDeployerIT.java | 65 +++++++++++++++++++
 .../autoredeploy/TracedSmallModelRequest.json | 14 ++++
 2 files changed, 79 insertions(+)
 create mode 100644 plugin/src/test/java/org/opensearch/ml/autoredeploy/MLModelAutoReDeployerIT.java
 create mode 100644 plugin/src/test/resources/org/opensearch/ml/autoredeploy/TracedSmallModelRequest.json

diff --git a/plugin/src/test/java/org/opensearch/ml/autoredeploy/MLModelAutoReDeployerIT.java b/plugin/src/test/java/org/opensearch/ml/autoredeploy/MLModelAutoReDeployerIT.java
new file mode 100644
index 0000000000..1ab8218d2e
--- /dev/null
+++ b/plugin/src/test/java/org/opensearch/ml/autoredeploy/MLModelAutoReDeployerIT.java
@@ -0,0 +1,65 @@
+/*
+ *
+ *  * Copyright OpenSearch Contributors
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package org.opensearch.ml.autoredeploy;
+
+import static org.opensearch.ml.common.MLTask.MODEL_ID_FIELD;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+
+import org.opensearch.ml.common.MLTaskState;
+import org.opensearch.ml.rest.MLCommonsRestTestCase;
+
+import lombok.SneakyThrows;
+
+public class MLModelAutoReDeployerIT extends MLCommonsRestTestCase {
+
+    public void testModelAutoRedeploy() {
+        prepareModel();
+    }
+
+    @SneakyThrows
+    private void prepareModel() {
+        String requestBody = Files
+            .readString(
+                Path.of(this.getClass().getClassLoader().getResource("org/opensearch/ml/autoredeploy/TracedSmallModelRequest.json").toURI())
+            );
+        String registerFirstModelTaskId = registerModel(requestBody);
+        String registerSecondModelTaskId = registerModel(requestBody);
+        waitForTask(registerFirstModelTaskId, MLTaskState.COMPLETED);
+        getTask(client(), registerFirstModelTaskId, response -> {
+            String firstModelId = (String) response.get(MODEL_ID_FIELD);
+            try {
+                String deployFirstModelTaskId = deployModel(firstModelId);
+                getTask(client(), registerSecondModelTaskId, response1 -> {
+                    String secondModelId = (String) response1.get(MODEL_ID_FIELD);
+                    try {
+                        /**
+                         * At this time point, the model auto redeployer should be querying the deploying/deploy failed/partially deployed models.
+                         * The original deploy model task should be able to complete successfully, if not it means the
+                         * org.opensearch.ml.action.forward.TransportForwardAction.triggerNextModelDeployAndCheckIfRestRetryTimes might throw exception
+                         * which cause by org.opensearch.ml.autoredeploy.MLModelAutoReDeployer#redeployAModel. The auto redeploy constructs an arrangement
+                         * with two models, the first model deploy done event will trigger the auto redeploy's next model deploy, and if during this
+                         * any error occurs, the first model deploy task status won't be updated to complete. So if this IT can pass, then it means the
+                         * next model auto redeploy trigger is correct.
+                         */
+                        String deploySecondModelTaskId = deployModel(secondModelId);
+                        waitForTask(deploySecondModelTaskId, MLTaskState.COMPLETED);
+                    } catch (Exception e) {
+                        fail(e.getMessage());
+                    }
+                });
+                waitForTask(deployFirstModelTaskId, MLTaskState.COMPLETED);
+            } catch (Exception e) {
+                logger.error(e.getMessage(), e);
+                fail(e.getMessage());
+            }
+        });
+    }
+
+}
diff --git a/plugin/src/test/resources/org/opensearch/ml/autoredeploy/TracedSmallModelRequest.json b/plugin/src/test/resources/org/opensearch/ml/autoredeploy/TracedSmallModelRequest.json
new file mode 100644
index 0000000000..0173665a5d
--- /dev/null
+++ b/plugin/src/test/resources/org/opensearch/ml/autoredeploy/TracedSmallModelRequest.json
@@ -0,0 +1,14 @@
+{
+  "name": "traced_small_model",
+  "version": "1.0.0",
+  "model_format": "TORCH_SCRIPT",
+  "model_task_type": "text_embedding",
+  "model_content_hash_value": "e13b74006290a9d0f58c1376f9629d4ebc05a0f9385f40db837452b167ae9021",
+  "model_config": {
+    "model_type": "bert",
+    "embedding_dimension": 768,
+    "framework_type": "sentence_transformers",
+    "all_config": "{\"architectures\":[\"BertModel\"],\"max_position_embeddings\":512,\"model_type\":\"bert\",\"num_attention_heads\":12,\"num_hidden_layers\":6}"
+  },
+  "url": "https://github.com/opensearch-project/ml-commons/blob/2.x/ml-algorithms/src/test/resources/org/opensearch/ml/engine/algorithms/text_embedding/traced_small_model.zip?raw=true"
+}
\ No newline at end of file