From 4ff68c194cbb34e01e5d42873416ec8e10ce739d Mon Sep 17 00:00:00 2001
From: Enrico Olivelli <enrico.olivelli@datastax.com>
Date: Thu, 19 Oct 2023 04:46:49 -0400
Subject: [PATCH] [query agent] Add test case with output-field that contains a
 dash character (#616)

---
 .../datastax/oss/streaming/ai/QueryStep.java  | 10 +++---
 .../oss/streaming/ai/QueryStepTest.java       | 36 +++++++++++++++++++
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/QueryStep.java b/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/QueryStep.java
index a4b45a2e4..ce4d0b631 100644
--- a/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/QueryStep.java
+++ b/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/QueryStep.java
@@ -70,10 +70,12 @@ public class QueryStep implements TransformStep {
         this.dataSource = dataSource;
         this.loopOver = loopOver;
         this.loopOverAccessor = loopOverAccessor;
-        this.fields.forEach(
-                field -> {
-                    fieldsEvaluators.add(new JstlEvaluator<>("${" + field + "}", Object.class));
-                });
+        if (this.fields != null) {
+            this.fields.forEach(
+                    field -> {
+                        fieldsEvaluators.add(new JstlEvaluator<>("${" + field + "}", Object.class));
+                    });
+        }
         if (loopOver != null && !loopOver.isEmpty()) {
             this.loopOverAccessor = new JstlEvaluator<>("${" + loopOver + "}", List.class);
         } else {
diff --git a/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/QueryStepTest.java b/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/QueryStepTest.java
index c2d6b60fa..890bdfd38 100644
--- a/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/QueryStepTest.java
+++ b/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/QueryStepTest.java
@@ -552,4 +552,40 @@ public List<Map<String, Object>> fetchData(String query, List<Object> params) {
                 (List<Map<String, Object>>) result.get("command_results");
         assertEquals(List.of(Map.of("foo", "bar"), Map.of("foo", "bar2")), command_results);
     }
+
+    @Test
+    void testSetFieldWithDash() throws Exception {
+
+        String value = "{}";
+
+        QueryStepDataSource dataSource =
+                new QueryStepDataSource() {
+
+                    @Override
+                    public List<Map<String, Object>> fetchData(String query, List<Object> params) {
+                        return List.of(Map.of("foo", "bar"));
+                    }
+                };
+
+        QueryStep queryStep =
+                QueryStep.builder()
+                        .dataSource(dataSource)
+                        .outputFieldName("value.command-results")
+                        .query("select 1")
+                        .build();
+        queryStep.start();
+
+        MutableRecord context =
+                MutableRecord.recordToMutableRecord(SimpleRecord.of(null, value), true);
+
+        queryStep.process(context);
+        ai.langstream.api.runner.code.Record record =
+                MutableRecord.mutableRecordToRecord(context).orElseThrow();
+        Map<String, Object> result = (Map<String, Object>) record.value();
+        log.info("Result: {}", result);
+
+        List<Map<String, Object>> command_results =
+                (List<Map<String, Object>>) result.get("command-results");
+        assertEquals(List.of(Map.of("foo", "bar")), command_results);
+    }
 }