20
20
import org .opensearch .core .rest .RestStatus ;
21
21
import org .opensearch .ml .common .MLModel ;
22
22
import org .opensearch .ml .common .model .MLModelState ;
23
+ import org .opensearch .ml .common .utils .StringUtils ;
23
24
import org .opensearch .ml .rest .RestBaseAgentToolsIT ;
24
25
import org .opensearch .ml .utils .TestHelper ;
25
26
@@ -77,9 +78,9 @@ public void deleteModel() throws IOException {
77
78
78
79
@ SneakyThrows
79
80
private void checkForModelUndeployedStatus (String modelId ) {
80
- Predicate <Response > condition = response -> {
81
+ Predicate <String > condition = responseStr -> {
81
82
try {
82
- Map <String , Object > responseInMap = parseResponseToMap ( response );
83
+ Map <String , Object > responseInMap = StringUtils . gson . fromJson ( responseStr , Map . class );
83
84
MLModelState state = MLModelState .from (responseInMap .get (MLModel .MODEL_STATE_FIELD ).toString ());
84
85
return MLModelState .UNDEPLOYED .equals (state );
85
86
} catch (Exception e ) {
@@ -90,14 +91,23 @@ private void checkForModelUndeployedStatus(String modelId) {
90
91
}
91
92
92
93
@ SneakyThrows
93
- protected Response waitResponseMeetingCondition (String method , String endpoint , String jsonEntity , Predicate <Response > condition ) {
94
+ protected Response waitResponseMeetingCondition (String method , String endpoint , String jsonEntity , Predicate <String > condition ) {
94
95
for (int attempt = 1 ; attempt <= MAX_RETRIES ; attempt ++) {
95
96
Response response = TestHelper .makeRequest (client (), method , endpoint , null , jsonEntity , null );
96
97
assertEquals (RestStatus .OK , RestStatus .fromCode (response .getStatusLine ().getStatusCode ()));
97
- if (condition .test (response )) {
98
+ String entityString = TestHelper .httpEntityToString (response .getEntity ());
99
+ if (condition .test (entityString )) {
98
100
return response ;
99
101
}
100
- logger .info ("The {}-th attempt on {}:{} . response: {}" , attempt , method , endpoint , response .toString ());
102
+ logger
103
+ .info (
104
+ "The {}-th attempt on {}:{} . response: {}, responseBody: {}" ,
105
+ attempt ,
106
+ method ,
107
+ endpoint ,
108
+ response .toString (),
109
+ entityString
110
+ );
101
111
Thread .sleep (DEFAULT_TASK_RESULT_QUERY_INTERVAL_IN_MILLISECOND );
102
112
}
103
113
fail (
0 commit comments