From 48ee20782fdc73bf4b28a68e7c0de531842a6568 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20N=2E=20Robalino?= <andres@androbtech.com>
Date: Fri, 29 May 2020 03:10:15 -0500
Subject: [PATCH] Ensure `end_filter` plugin lifecycle stage gets called.

---
 crates/nu-cli/src/commands/plugin.rs  | 79 ++++++++++++++++-----------
 docs/commands/average.md              |  2 +-
 tests/plugins/decommission_average.rs | 26 +++++++++
 tests/plugins/mod.rs                  |  1 +
 4 files changed, 74 insertions(+), 34 deletions(-)
 create mode 100644 tests/plugins/decommission_average.rs

diff --git a/crates/nu-cli/src/commands/plugin.rs b/crates/nu-cli/src/commands/plugin.rs
index 0d231d9cfc..e4ace167e0 100644
--- a/crates/nu-cli/src/commands/plugin.rs
+++ b/crates/nu-cli/src/commands/plugin.rs
@@ -194,7 +194,7 @@ pub fn filter_plugin(
             }
         }
 
-        // End of the stream
+        // post stream contents
         {
             let stdin = child.stdin.as_mut().expect("Failed to open stdin");
             let stdout = child.stdout.as_mut().expect("Failed to open stdout");
@@ -202,15 +202,23 @@ pub fn filter_plugin(
             let mut reader = BufReader::new(stdout);
 
             let request: JsonRpc<std::vec::Vec<Value>> = JsonRpc::new("end_filter", vec![]);
-            let request_raw = match serde_json::to_string(&request) {
-                Ok(req) => req,
-                Err(err) => {
-                    yield Err(ShellError::unexpected(format!("{}", err)));
-                    return;
-                }
-            };
+            let request_raw = serde_json::to_string(&request);
 
-            let _ = stdin.write(format!("{}\n", request_raw).as_bytes()); // TODO: Handle error
+            match request_raw {
+                Err(_) => {
+                    yield Err(ShellError::labeled_error(
+                        "Could not load json from plugin",
+                        "could not load json from plugin",
+                        &call_info.name_tag,
+                    ));
+                }
+                Ok(request_raw) => match stdin.write(format!("{}\n", request_raw).as_bytes()) {
+                    Ok(_) => {}
+                    Err(err) => {
+                        yield Err(ShellError::unexpected(format!("{}", err)));
+                    }
+                },
+            }
 
             let mut input = String::new();
             match reader.read_line(&mut input) {
@@ -218,26 +226,7 @@ pub fn filter_plugin(
                     let response = serde_json::from_str::<NuResult>(&input);
                     match response {
                         Ok(NuResult::response { params }) => match params {
-                            Ok(params) => {
-                                let request: JsonRpc<std::vec::Vec<Value>> =
-                                    JsonRpc::new("quit", vec![]);
-                                let request_raw = serde_json::to_string(&request);
-                                match request_raw {
-                                    Ok(request_raw) => {
-                                        let _ = stdin.write(format!("{}\n", request_raw).as_bytes()); // TODO: Handle error
-                                    }
-                                    Err(e) => {
-                                        yield Err(ShellError::untagged_runtime_error(format!(
-                                            "Error while processing begin_filter response: {:?} {}",
-                                            e, input
-                                        )));
-                                        return;
-                                    }
-                                }
-
-                                //yield ReturnValue::Ok(params)
-                                //yield ReturnSuccess::value(Value)
-                            }
+                            Ok(params) => for param in params { yield param },
                             Err(e) => {
                                 yield ReturnValue::Err(e);
                             }
@@ -252,14 +241,38 @@ pub fn filter_plugin(
                 }
                 Err(e) => {
                     yield Err(ShellError::untagged_runtime_error(format!(
-                        "Error while reading end_filter: {:?}",
+                        "Error while reading end_filter response: {:?}",
                         e
                     )));
                 }
-            };
-
-            let _ = child.wait();
+            }
         }
+
+        // End of the stream
+        {
+            let stdin = child.stdin.as_mut().expect("Failed to open stdin");
+            let stdout = child.stdout.as_mut().expect("Failed to open stdout");
+
+            let mut reader = BufReader::new(stdout);
+
+            let request: JsonRpc<std::vec::Vec<Value>> = JsonRpc::new("quit", vec![]);
+            let request_raw = serde_json::to_string(&request);
+
+            match request_raw {
+                Ok(request_raw) => {
+                    let _ = stdin.write(format!("{}\n", request_raw).as_bytes()); // TODO: Handle error
+                }
+                Err(e) => {
+                    yield Err(ShellError::untagged_runtime_error(format!(
+                        "Error while processing quit response: {:?}",
+                        e
+                    )));
+                    return;
+                }
+            }
+        }
+
+        let _ = child.wait();
     };
 
     Ok(stream.to_output_stream())
diff --git a/docs/commands/average.md b/docs/commands/average.md
index c2072791d9..3923ae9e54 100644
--- a/docs/commands/average.md
+++ b/docs/commands/average.md
@@ -1,5 +1,5 @@
 # average
-This command allows you to calculate the average of values in a column.
+Calculate the average of values in a column.
 
 ## Examples
 To get the average of the file sizes in a directory, simply pipe the size column from the ls command to the average command.
diff --git a/tests/plugins/decommission_average.rs b/tests/plugins/decommission_average.rs
new file mode 100644
index 0000000000..1601076171
--- /dev/null
+++ b/tests/plugins/decommission_average.rs
@@ -0,0 +1,26 @@
+use nu_test_support::{nu, pipeline};
+
+#[test]
+fn can_average_numbers() {
+    let actual = nu!(
+        cwd: "tests/fixtures/formats", pipeline(
+        r#"
+             open sgml_description.json
+             | get glossary.GlossDiv.GlossList.GlossEntry.Sections
+             | average
+             | echo $it
+         "#
+    ));
+
+    assert_eq!(actual.out, "101.5")
+}
+
+#[test]
+fn can_average_bytes() {
+    let actual = nu!(
+        cwd: "tests/fixtures/formats",
+        "ls | sort-by name | skip 1 | first 2 | get size | average | format \"{$it}\" | echo $it"
+    );
+
+    assert_eq!(actual.out, "1.6 KB");
+}
diff --git a/tests/plugins/mod.rs b/tests/plugins/mod.rs
index e3d59b2205..7add8f932b 100644
--- a/tests/plugins/mod.rs
+++ b/tests/plugins/mod.rs
@@ -1 +1,2 @@
 mod core_inc;
+mod decommission_average;