Fixing register agent tests

Signed-off-by: Joshua Palis <[email protected]>
opensearch-project · Mar 19, 2024 · b4fa72e · b4fa72e
1 parent aa8a454
commit b4fa72e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/test/java/org/opensearch/flowframework/workflow/RegisterAgentTests.java b/src/test/java/org/opensearch/flowframework/workflow/RegisterAgentTests.java
@@ -16,6 +16,7 @@
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.indices.FlowFrameworkIndicesHandler;
 import org.opensearch.ml.client.MachineLearningNodeClient;
+import org.opensearch.ml.common.MLAgentType;
 import org.opensearch.ml.common.agent.LLMSpec;
 import org.opensearch.ml.common.agent.MLAgent;
 import org.opensearch.ml.common.agent.MLMemorySpec;
@@ -69,7 +70,7 @@ public void setUp() throws Exception {
             Map.ofEntries(
                 Map.entry("name", "test"),
                 Map.entry("description", "description"),
-                Map.entry("type", "type"),
+                Map.entry("type", MLAgentType.FLOW.name()),
                 Map.entry("llm.model_id", "xyz"),
                 Map.entry("llm.parameters", Collections.emptyMap()),
                 Map.entry("tools", tools),