From 6d9f725663f298587b53334838ea5079c4dfc887 Mon Sep 17 00:00:00 2001
From: Hailong Cui <ihailong@amazon.com>
Date: Thu, 13 Feb 2025 18:12:39 +0800
Subject: [PATCH] add more logs to toubleshot flaky test

Signed-off-by: Hailong Cui <ihailong@amazon.com>
---
 .../ml/tools/ToolIntegrationWithLLMTest.java  | 20 ++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/plugin/src/test/java/org/opensearch/ml/tools/ToolIntegrationWithLLMTest.java b/plugin/src/test/java/org/opensearch/ml/tools/ToolIntegrationWithLLMTest.java
index a9b3f84bfc..f73478d092 100644
--- a/plugin/src/test/java/org/opensearch/ml/tools/ToolIntegrationWithLLMTest.java
+++ b/plugin/src/test/java/org/opensearch/ml/tools/ToolIntegrationWithLLMTest.java
@@ -20,6 +20,7 @@
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.ml.common.MLModel;
 import org.opensearch.ml.common.model.MLModelState;
+import org.opensearch.ml.common.utils.StringUtils;
 import org.opensearch.ml.rest.RestBaseAgentToolsIT;
 import org.opensearch.ml.utils.TestHelper;
 
@@ -77,9 +78,9 @@ public void deleteModel() throws IOException {
 
     @SneakyThrows
     private void checkForModelUndeployedStatus(String modelId) {
-        Predicate<Response> condition = response -> {
+        Predicate<String> condition = responseStr -> {
             try {
-                Map<String, Object> responseInMap = parseResponseToMap(response);
+                Map<String, Object> responseInMap = StringUtils.gson.fromJson(responseStr, Map.class);
                 MLModelState state = MLModelState.from(responseInMap.get(MLModel.MODEL_STATE_FIELD).toString());
                 return MLModelState.UNDEPLOYED.equals(state);
             } catch (Exception e) {
@@ -90,14 +91,23 @@ private void checkForModelUndeployedStatus(String modelId) {
     }
 
     @SneakyThrows
-    protected Response waitResponseMeetingCondition(String method, String endpoint, String jsonEntity, Predicate<Response> condition) {
+    protected Response waitResponseMeetingCondition(String method, String endpoint, String jsonEntity, Predicate<String> condition) {
         for (int attempt = 1; attempt <= MAX_RETRIES; attempt++) {
             Response response = TestHelper.makeRequest(client(), method, endpoint, null, jsonEntity, null);
             assertEquals(RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode()));
-            if (condition.test(response)) {
+            String entityString = TestHelper.httpEntityToString(response.getEntity());
+            if (condition.test(entityString)) {
                 return response;
             }
-            logger.info("The {}-th attempt on {}:{} . response: {}", attempt, method, endpoint, response.toString());
+            logger
+                .info(
+                    "The {}-th attempt on {}:{} . response: {}, responseBody: {}",
+                    attempt,
+                    method,
+                    endpoint,
+                    response.toString(),
+                    entityString
+                );
             Thread.sleep(DEFAULT_TASK_RESULT_QUERY_INTERVAL_IN_MILLISECOND);
         }
         fail(