add more complicated test cases

Signed-off-by: Mingshi Liu <[email protected]>
opensearch-project · Aug 14, 2024 · 99d3796 · 99d3796
1 parent 83572b5
commit 99d3796
Show file tree

Hide file tree

Showing 2 changed files with 90 additions and 6 deletions.
diff --git a/common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java b/common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java
@@ -324,16 +324,15 @@ public <T> T createPayload(String action, Map<String, String> parameters) {
  payload = fillNullParameters(parameters, payload);
  StringSubstitutor substitutor = new StringSubstitutor(parameters, "${parameters.", "}");
  payload = substitutor.replace(payload);
- boolean isJson = isJson(payload);
- if (!isJson) {
- String manuallyFixedJson = connectorAction.get().getRequestBody();
+ if (!isJson(payload)) {
+ String payloadAfterEscape = connectorAction.get().getRequestBody();
  Map<String, String> escapedParameters = escapeMapValues(parameters);
  StringSubstitutor escapedSubstitutor = new StringSubstitutor(escapedParameters, "${parameters.", "}");
- manuallyFixedJson = escapedSubstitutor.replace(manuallyFixedJson);
- if (!isJson(manuallyFixedJson)) {
+ payloadAfterEscape = escapedSubstitutor.replace(payloadAfterEscape);
+ if (!isJson(payloadAfterEscape)) {
  throw new IllegalArgumentException("Invalid payload: " + payload);
  } else {
- payload = manuallyFixedJson;
+ payload = payloadAfterEscape;
  }
  }
  return (T) payload;

diff --git a/common/src/test/java/org/opensearch/ml/common/connector/HttpConnectorTest.java b/common/src/test/java/org/opensearch/ml/common/connector/HttpConnectorTest.java
@@ -197,6 +197,10 @@ public void createPayloadWithString() {
  Assert.assertEquals("{\"prompt\": \"answer question based on context: document1\"}", predictPayload);
  }
 
+ /**
+ *
+ */
+
  @Test
  public void createPayloadWithList() {
  String requestBody = "{\"prompt\": \"${parameters.prompt}\"}";
@@ -211,6 +215,87 @@ public void createPayloadWithList() {
  connector.validatePayload(predictPayload);
  }
 
+ @Test
+ public void createPayloadWithNestedList() {
+ String requestBody = "{\"prompt\": \"${parameters.prompt}\"}";
+ HttpConnector connector = createHttpConnectorWithRequestBody(requestBody);
+ Map<String, String> parameters = new HashMap<>();
+ parameters.put("prompt", "answer question based on context: ${parameters.context}");
+ ArrayList<String> listOfDocuments = new ArrayList<>();
+ listOfDocuments.add("document1");
+ ArrayList<String> NestedListOfDocuments = new ArrayList<>();
+ NestedListOfDocuments.add("document2");
+ listOfDocuments.add(toJson(NestedListOfDocuments));
+ parameters.put("context", toJson(listOfDocuments));
+ String predictPayload = connector.createPayload(PREDICT.name(), parameters);
+ connector.validatePayload(predictPayload);
+ }
+
+ @Test
+ public void createPayloadWithMap() {
+ String requestBody = "{\"prompt\": \"${parameters.prompt}\"}";
+ HttpConnector connector = createHttpConnectorWithRequestBody(requestBody);
+ Map<String, String> parameters = new HashMap<>();
+ parameters.put("prompt", "answer question based on context: ${parameters.context}");
+ Map<String, String> mapOfDocuments = new HashMap<>();
+ mapOfDocuments.put("name", "John");
+ parameters.put("context", toJson(mapOfDocuments));
+ String predictPayload = connector.createPayload(PREDICT.name(), parameters);
+ connector.validatePayload(predictPayload);
+ }
+
+ @Test
+ public void createPayloadWithNestedMapOfString() {
+ String requestBody = "{\"prompt\": \"${parameters.prompt}\"}";
+ HttpConnector connector = createHttpConnectorWithRequestBody(requestBody);
+ Map<String, String> parameters = new HashMap<>();
+ parameters.put("prompt", "answer question based on context: ${parameters.context}");
+ Map<String, String> mapOfDocuments = new HashMap<>();
+ mapOfDocuments.put("name", "John");
+ Map<String, String> nestedMapOfDocuments = new HashMap<>();
+ nestedMapOfDocuments.put("city", "New York");
+ mapOfDocuments.put("hometown", toJson(nestedMapOfDocuments));
+ parameters.put("context", toJson(mapOfDocuments));
+ String predictPayload = connector.createPayload(PREDICT.name(), parameters);
+ connector.validatePayload(predictPayload);
+ }
+
+ @Test
+ public void createPayloadWithNestedMapOfObject() {
+ String requestBody = "{\"prompt\": \"${parameters.prompt}\"}";
+ HttpConnector connector = createHttpConnectorWithRequestBody(requestBody);
+ Map<String, String> parameters = new HashMap<>();
+ parameters.put("prompt", "answer question based on context: ${parameters.context}");
+ Map<String, Object> mapOfDocuments = new HashMap<>();
+ mapOfDocuments.put("name", "John");
+ Map<String, String> nestedMapOfDocuments = new HashMap<>();
+ nestedMapOfDocuments.put("city", "New York");
+ mapOfDocuments.put("hometown", nestedMapOfDocuments);
+ parameters.put("context", toJson(mapOfDocuments));
+ String predictPayload = connector.createPayload(PREDICT.name(), parameters);
+ connector.validatePayload(predictPayload);
+ }
+
+ @Test
+ public void createPayloadWithNestedListOfMapOfObject() {
+ String requestBody = "{\"prompt\": \"${parameters.prompt}\"}";
+ HttpConnector connector = createHttpConnectorWithRequestBody(requestBody);
+ Map<String, String> parameters = new HashMap<>();
+ parameters.put("prompt", "answer question based on context: ${parameters.context}");
+ ArrayList<String> listOfDocuments = new ArrayList<>();
+ listOfDocuments.add("document1");
+ ArrayList<Object> NestedListOfDocuments = new ArrayList<>();
+ Map<String, Object> mapOfDocuments = new HashMap<>();
+ mapOfDocuments.put("name", "John");
+ Map<String, String> nestedMapOfDocuments = new HashMap<>();
+ nestedMapOfDocuments.put("city", "New York");
+ mapOfDocuments.put("hometown", nestedMapOfDocuments);
+ listOfDocuments.add(toJson(NestedListOfDocuments));
+ parameters.put("context", toJson(listOfDocuments));
+ String predictPayload = connector.createPayload(PREDICT.name(), parameters);
+ connector.validatePayload(predictPayload);
+ }
+
  @Test
  public void createPayload() {
  HttpConnector connector = createHttpConnector();