From 04faf76f6038644778a6dd6fad08157d2ff7cd38 Mon Sep 17 00:00:00 2001
From: anussel5559 <alex@redoxengine.com>
Date: Sat, 7 Sep 2024 09:22:15 -0500
Subject: [PATCH] add-select: adjust dynamic selector shape and render flow,
 tests

---
 crates/slumber_core/src/collection/models.rs  |   9 +-
 crates/slumber_core/src/http/query.rs         |   9 +-
 crates/slumber_core/src/template.rs           | 124 ++++++++++++++++--
 crates/slumber_core/src/template/error.rs     |   5 +
 crates/slumber_core/src/template/render.rs    |  50 +++++--
 .../src/view/component/queryable_body.rs      |   2 +-
 slumber.yml                                   |   7 +-
 7 files changed, 173 insertions(+), 33 deletions(-)

diff --git a/crates/slumber_core/src/collection/models.rs b/crates/slumber_core/src/collection/models.rs
index c4a97fe3..58bdbf16 100644
--- a/crates/slumber_core/src/collection/models.rs
+++ b/crates/slumber_core/src/collection/models.rs
@@ -435,8 +435,13 @@ pub enum ChainSource {
 #[serde(rename_all = "snake_case", deny_unknown_fields, untagged)]
 pub enum SelectOptions {
     Fixed(Vec<Template>),
-    /// Template will be rendered as a parsed JSON array to get options
-    Dynamic(Template),
+    /// Dynamic requires a source (often a chain) that either returns a JSON
+    /// array OR a JSON object, in which case we'll use selector to query
+    /// and find the array to parse in to the list of options
+    Dynamic {
+        source: Template,
+        selector: Option<Query>,
+    },
 }
 
 /// Test-only helpers
diff --git a/crates/slumber_core/src/http/query.rs b/crates/slumber_core/src/http/query.rs
index 62ec498f..9909edce 100644
--- a/crates/slumber_core/src/http/query.rs
+++ b/crates/slumber_core/src/http/query.rs
@@ -3,7 +3,7 @@
 use crate::http::content_type::ResponseContent;
 use derive_more::{Display, FromStr};
 use serde::{Deserialize, Serialize};
-use serde_json_path::{ExactlyOneError, JsonPath};
+use serde_json_path::{ExactlyOneError, JsonPath, NodeList};
 use std::borrow::Cow;
 use thiserror::Error;
 
@@ -17,7 +17,7 @@ impl Query {
     /// Apply a query to some content, returning the result in the original
     /// format. This will convert to a common format, apply the query, then
     /// convert back.
-    pub fn query(
+    pub fn query_content(
         &self,
         value: &dyn ResponseContent,
     ) -> Box<dyn ResponseContent> {
@@ -30,6 +30,11 @@ impl Query {
         content_type.parse_json(Cow::Owned(queried))
     }
 
+    /// Apply a query to some content, returning a list of results.
+    pub fn query_list<'a>(&self, value: &'a serde_json::Value) -> NodeList<'a> {
+        self.0.query(value)
+    }
+
     /// Apply a query to some content, returning a string. The query should
     /// return a single result. If it's a scalar, that will be stringified. If
     /// it's an array/object, it'll be converted back into its input format,
diff --git a/crates/slumber_core/src/template.rs b/crates/slumber_core/src/template.rs
index feb53052..c8be1aaa 100644
--- a/crates/slumber_core/src/template.rs
+++ b/crates/slumber_core/src/template.rs
@@ -276,7 +276,7 @@ mod tests {
     use chrono::Utc;
     use indexmap::indexmap;
     use rstest::rstest;
-    use serde_json::json;
+    use serde_json::{json, Value};
     use std::time::Duration;
     use tokio::fs;
     use wiremock::{matchers, Mock, MockServer, ResponseTemplate};
@@ -947,12 +947,12 @@ mod tests {
     }
 
     #[rstest]
-    #[case::no_chains(vec!["foo!", "bar!"], 0, "foo!")]
-    #[case::chains_0(vec!["foo!", "{{chains.command}}"], 0, "foo!")]
-    #[case::chains_1(vec!["foo!", "{{chains.command}}"], 1, "command_output")]
+    #[case::no_chains(SelectOptions::Fixed(vec!["foo!".into(), "bar!".into()]), 0, "foo!")]
+    #[case::chains_0(SelectOptions::Fixed(vec!["foo!".into(), "{{chains.command}}".into()]), 0, "foo!")]
+    #[case::chains_1(SelectOptions::Fixed(vec!["foo!".into(), "{{chains.command}}".into()]), 1, "command_output")]
     #[tokio::test]
-    async fn test_chain_select(
-        #[case] options: Vec<&str>,
+    async fn test_chain_fixed_select(
+        #[case] options: SelectOptions,
         #[case] index: usize,
         #[case] expected: &str,
     ) {
@@ -960,11 +960,7 @@ mod tests {
             id: "sut".into(),
             source: ChainSource::Select {
                 message: Some("password".into()),
-                options: options
-                    .clone()
-                    .into_iter()
-                    .map(|s| s.into())
-                    .collect(),
+                options,
             },
             ..Chain::factory(())
         };
@@ -989,6 +985,71 @@ mod tests {
         assert_eq!(render!("{{chains.sut}}", context).unwrap(), expected);
     }
 
+    #[rstest]
+    #[case::dynamic_select_selector_0(
+        SelectOptions::Dynamic { source: "{{chains.command}}".into(), selector: Some("$.array[*]".parse().unwrap())},
+        json!({"array": ["foo", "bar"]}),
+        0,
+        "foo"
+    )]
+    #[case::dynamic_select_selector_1(
+        SelectOptions::Dynamic { source: "{{chains.command}}".into(), selector: Some("$.array[*]".parse().unwrap())},
+        json!({"array": ["foo", "bar"]}),
+        1,
+        "bar"
+    )]
+    #[case::dynamic_select_0(
+        SelectOptions::Dynamic { source: "{{chains.command}}".into(), selector: None},
+        json!(["foo", "bar"]),
+        0,
+        "foo"
+    )]
+    #[case::dynamic_select_1(
+        SelectOptions::Dynamic { source: "{{chains.command}}".into(), selector: None},
+        json!(["foo", "bar"]),
+        1,
+        "bar"
+    )]
+    #[tokio::test]
+    async fn test_chain_dynamic_select(
+        #[case] options: SelectOptions,
+        #[case] chain_json: Value,
+        #[case] index: usize,
+        #[case] expected: &str,
+    ) {
+        let sut_chain = Chain {
+            id: "sut".into(),
+            source: ChainSource::Select {
+                message: Some("password".into()),
+                options,
+            },
+            ..Chain::factory(())
+        };
+
+        let command_chain = Chain {
+            id: "command".into(),
+            source: ChainSource::command([
+                "echo",
+                "-n",
+                format!("{}", chain_json).as_str(),
+            ]),
+            trim: ChainOutputTrim::Both,
+            ..Chain::factory(())
+        };
+
+        let context = TemplateContext {
+            collection: Collection {
+                chains: by_id([sut_chain, command_chain]),
+                ..Collection::factory(())
+            }
+            .into(),
+            prompter: Box::new(TestSelectPrompter::new(vec![index])),
+            ..TemplateContext::factory(())
+        };
+
+        assert_eq!(render!("{{chains.sut}}", context).unwrap(), expected);
+    }
+
     #[tokio::test]
     async fn test_chain_select_error() {
         let chain = Chain {
@@ -1015,6 +1076,47 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_chain_select_dynamic_error() {
+        let sut_chain = Chain {
+            source: ChainSource::Select {
+                message: Some("password".into()),
+                options: SelectOptions::Dynamic {
+                    source: "{{chains.command}}".into(),
+                    selector: None,
+                },
+            },
+            ..Chain::factory(())
+        };
+
+        let chain_json = json!("{ not_an_array: 42 }");
+
+        let command_chain = Chain {
+            id: "command".into(),
+            source: ChainSource::command([
+                "echo",
+                format!("{}", chain_json).as_str(),
+            ]),
+            trim: ChainOutputTrim::Both,
+            ..Chain::factory(())
+        };
+
+        let context = TemplateContext {
+            collection: Collection {
+                chains: by_id([sut_chain, command_chain]),
+                ..Collection::factory(())
+            }
+            .into(),
+            prompter: Box::new(TestSelectPrompter::new(vec![0 as usize])),
+            ..TemplateContext::factory(())
+        };
+
+        assert_err!(
+            render!("{{chains.chain1}}", context),
+            "Dynamic select options are invalid. Source must be an array, or a selector must be provided."
+        );
+    }
+
     /// Test that a chain being used twice only computes the chain once
     #[tokio::test]
     async fn test_chain_duplicate() {
diff --git a/crates/slumber_core/src/template/error.rs b/crates/slumber_core/src/template/error.rs
index a3895792..8c9c2d5f 100644
--- a/crates/slumber_core/src/template/error.rs
+++ b/crates/slumber_core/src/template/error.rs
@@ -178,12 +178,17 @@ pub enum ChainError {
     #[error("No response from prompt/select")]
     PromptNoResponse,
 
+    /// We hit some sort of deserialization error while trying
+    /// to build dynamic options
     #[error("Dynamic list failed to deserialize to a value array")]
     DynamicSelectOptions {
         #[source]
         error: Arc<serde_json::Error>,
     },
 
+    #[error("Dynamic select options are invalid. Source must be an array, or a selector must be provided.")]
+    DynamicOptionsInvalid,
+
     /// A bubbled-error from rendering a nested template in the chain arguments
     #[error("Rendering nested template for field `{field}`")]
     Nested {
diff --git a/crates/slumber_core/src/template/render.rs b/crates/slumber_core/src/template/render.rs
index d8e7cfc2..414f5e9f 100644
--- a/crates/slumber_core/src/template/render.rs
+++ b/crates/slumber_core/src/template/render.rs
@@ -771,28 +771,50 @@ impl<'a> ChainTemplateSource<'a> {
                 ))
                 .await?
             }
-            SelectOptions::Dynamic(options) => {
-                let dynamic_options = options
-                    .render_chain_config("options", context, stack)
+            SelectOptions::Dynamic { source, selector } => {
+                // Render source to a string
+                let source = source
+                    .render_chain_config("source", context, stack)
                     .await?;
 
-                let raw_results: Vec<serde_json::Value> = serde_json::from_str(
-                    &dynamic_options,
-                )
-                .map_err(|error| ChainError::DynamicSelectOptions {
-                    error: Arc::new(error),
-                })?;
+                // the above render_chain_config will always parse the output of
+                // the chain as a string, but we'll need to
+                // parse back to JSON to execute the JSONPath
+                // query against it.
+                let source_json: serde_json::Value =
+                    serde_json::from_str(&source).map_err(|error| {
+                        ChainError::DynamicSelectOptions {
+                            error: Arc::new(error),
+                        }
+                    })?;
 
-                raw_results
+                let options = match (selector, source_json) {
+                    // A selector and raw JSON value means we execute
+                    // the selector
+                    (Some(selector), value) => Ok(selector
+                        .query_list(&value)
+                        .into_iter()
+                        .cloned()
+                        .collect()),
+                    // No selector, but a JSON array means we're good to go!
+                    (None, serde_json::Value::Array(options)) => Ok(options),
+
+                    (None, _) => Err(ChainError::DynamicOptionsInvalid),
+                }?;
+
+                // Convert all values to strings, for convenience. THis may be a
+                // bit wonky for nexted objects/arrays, but it
+                // should be obvious to the user what's going on so
+                // it's better than returning an error
+                options
                     .iter()
                     .map(|value| {
                         // if the value is already a string, avoid calling
                         // to_string on it since that would wrap the string
                         // in another set of quotes.
-                        if value.is_string() {
-                            value.as_str().unwrap().into()
-                        } else {
-                            value.to_string()
+                        match value {
+                            serde_json::Value::String(s) => s.clone(),
+                            other => other.to_string(),
                         }
                     })
                     .collect()
diff --git a/crates/slumber_tui/src/view/component/queryable_body.rs b/crates/slumber_tui/src/view/component/queryable_body.rs
index 191caad5..3bf1958a 100644
--- a/crates/slumber_tui/src/view/component/queryable_body.rs
+++ b/crates/slumber_tui/src/view/component/queryable_body.rs
@@ -232,7 +232,7 @@ fn init_text(
             // Body is a known content type so we parsed it - apply a query if
             // necessary and prettify the output
             query
-                .map(|query| query.query(parsed_body).prettify())
+                .map(|query| query.query_content(parsed_body).prettify())
                 .unwrap_or_else(|| parsed_body.prettify())
         })
         // Content couldn't be parsed, fall back to the raw text
diff --git a/slumber.yml b/slumber.yml
index df4c7c63..0ce74ed1 100644
--- a/slumber.yml
+++ b/slumber.yml
@@ -40,11 +40,12 @@ chains:
   dynamic_options:
     source: !request
       recipe: login
-    selector: $.args.fast
   select_dynamic:
     source: !select
       message: Select a value
-      options: "{{chains.dynamic_options}}"
+      options:
+        source: "{{chains.dynamic_options}}"
+        selector: $.args[*]
   auth_token:
     source: !request
       recipe: login
@@ -91,7 +92,7 @@ requests:
         url: "{{host}}/get"
         query:
           - foo=bar
-          - select={{chains.select_value}}
+          - select={{chains.select_dynamic}}
 
       get_user: !request
         <<: *base