2 files changed, 181 insertions, 19 deletions

diff --git a/src/hooks/wolfram_alpha.rs b/src/hooks/wolfram_alpha.rs
index e440eb8..5584935 100644
--- a/src/hooks/wolfram_alpha.rs
+++ b/src/hooks/wolfram_alpha.rs
@@ -1,4 +1,7 @@
-use crate::util::{formatting::truncate, web::shorten_url};
+use crate::util::{
+    formatting::truncate,
+    web::{quote_plus, shorten_url},
+};
 use anyhow::{bail, Context, Error, Result};
 use futures::try_join;
 use irc::client::prelude::*;
@@ -51,7 +54,7 @@ fn to_single_string(wa_res: WaResponse) -> String {
         .join(" - ")
 }
 
-fn get_url(query_str: &str, api_key: Option<&str>, base_url: Option<&str>) -> String {
+fn get_url(query_str: &str, api_key: Option<&str>, base_url: Option<&str>) -> Result<Url, Error> {
     let wa_url = "http://api.wolframalpha.com";
     let api_url = format!(
         "{}/v2/query?input={}&appid={}&output=json",
@@ -59,11 +62,11 @@ fn get_url(query_str: &str, api_key: Option<&str>, base_url: Option<&str>) -> St
         query_str,
         api_key.unwrap_or("XXX"), // Allow tests to run without a key
     );
-    api_url
+    Url::parse(&api_url).context("Failed to parse URL")
 }
 
-async fn send_wa_req(url: &str) -> Result<String, Error> {
-    let body = get(Url::parse(url).context("Failed to parse url")?)
+async fn send_wa_req(url: &Url) -> Result<String, Error> {
+    let body = get(url.to_owned())
         .await
         .context("Failed to make request")?
         .text()
@@ -72,12 +75,24 @@ async fn send_wa_req(url: &str) -> Result<String, Error> {
     Ok(body)
 }
 
-async fn handle_wa_req(url: &str) -> Result<WaResponse, Error> {
+async fn handle_wa_req(url: &Url) -> Result<WaResponse, Error> {
     let res_body = send_wa_req(url).await?;
     let parsed = serde_json::from_str(&res_body)?;
     Ok(parsed)
 }
 
+/// Gets a URL users can click, leading to the main WA page.
+async fn get_wa_user_short_url(input: &str) -> Result<String, Error> {
+    let user_url = format!(
+        "http://www.wolframalpha.com/input/?i={}",
+        // For some reason some inputs need double quote calls, e.g. '1 * 2'.
+        // Maybe only with is.gd though.
+        quote_plus(&quote_plus(input)?)?
+    );
+
+    shorten_url(&user_url).await
+}
+
 /// Sends a request to the Wolfram Alpha API, returns a plain text response.
 #[tracing::instrument]
 async fn wa_query(
@@ -85,34 +100,42 @@ async fn wa_query(
     api_key: Option<&str>,
     base_url: Option<&str>,
 ) -> Result<String, Error> {
-    let user_url = format!("http://www.wolframalpha.com/input/?i={}", query_str);
-    let user_url_shortened_fut = shorten_url(&user_url);
-
-    let url = get_url(query_str, api_key, base_url);
+    let user_url_shortened_fut = get_wa_user_short_url(query_str);
+    let url = get_url(query_str, api_key, base_url)?;
     let wa_res_fut = handle_wa_req(&url);
 
     // Can't just (foo.await, bar.await), smh
     // https://rust-lang.github.io/async-book/06_multiple_futures/02_join.html
     let (wa_res, user_url_shortened) = try_join!(wa_res_fut, user_url_shortened_fut)?;
 
+    let string_result = match to_single_string(wa_res) {
+        x if x.is_empty() => "No primary results.".to_string(),
+        x => x,
+    };
+
     Ok(format!(
         "{} - {}",
-        truncate(&to_single_string(wa_res), 250), // Same length as in gonzobot
+        truncate(&string_result, 250), // Same length as in gonzobot
         &user_url_shortened,
     ))
 }
 
+fn get_input_query(text: &str) -> Result<String, Error> {
+    let input = text.chars().as_str().splitn(2, " ").collect::<Vec<&str>>();
+    if input.len() != 2 {
+        bail!("Empty input for WA query");
+    }
+    let content = input[1].trim();
+    Ok(content.to_string())
+}
+
 pub async fn wa(bot: &crate::Bot, msg: Message) -> Result<()> {
     privmsg!(msg, {
-        let input = text.chars().as_str().splitn(2, " ").collect::<Vec<&str>>();
-        if input.len() != 2 {
-            bail!("Empty input for WA query");
-        }
-        let content = input[1].trim();
+        let content = get_input_query(text)?;
         bot.send_privmsg(
             msg.response_target()
                 .context("failed to get response target")?,
-            &wa_query(content, Some(&bot.config.settings.wa_api_key), None).await?,
+            &wa_query(&content, Some(&bot.config.settings.wa_api_key), None).await?,
         )?;
     })
 }
@@ -120,12 +143,86 @@ pub async fn wa(bot: &crate::Bot, msg: Message) -> Result<()> {
 #[cfg(test)]
 mod tests {
 
-    use super::wa_query;
+    use super::{get_input_query, get_wa_user_short_url, wa_query};
     use anyhow::{Error, Result};
     use mockito::{self, Matcher};
 
+    #[test]
+    fn test_input_query_content_retrieval() -> Result<(), Error> {
+        let incoming = ":wa test";
+        let content = get_input_query(incoming)?;
+        assert_eq!(content, "test");
+        Ok(())
+    }
+
+    #[test]
+    fn test_input_query_content_retrieval_with_spaces() -> Result<(), Error> {
+        let incoming = ":wa foo bar";
+        let content = get_input_query(incoming)?;
+        assert_eq!(content, "foo bar");
+        Ok(())
+    }
+
+    #[test]
+    fn test_input_query_content_retrieval_with_more_spaces() -> Result<(), Error> {
+        let incoming = ":wa foo bar baz";
+        let content = get_input_query(incoming)?;
+        assert_eq!(content, "foo bar baz");
+        Ok(())
+    }
+
+    // These tests must be updated if a service other than is.gd is used
     #[tokio::test]
-    async fn test_query_result_json_parsing() -> Result<(), Error> {
+    async fn test_wa_user_short_url_1() -> Result<(), Error> {
+        let input = "5/10";
+        assert_eq!(get_wa_user_short_url(input).await?, "https://is.gd/kgsSb7");
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_wa_user_short_url_2() -> Result<(), Error> {
+        let input = "5 / 10";
+        assert_eq!(get_wa_user_short_url(input).await?, "https://is.gd/b4xgwD");
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_wa_user_short_url_3() -> Result<(), Error> {
+        let input = "1*2";
+        assert_eq!(get_wa_user_short_url(input).await?, "https://is.gd/TLGKih");
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_wa_user_short_url_4() -> Result<(), Error> {
+        let input = "1 * 2";
+        assert_eq!(get_wa_user_short_url(input).await?, "https://is.gd/ZXYAUR");
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_wa_user_short_url_5() -> Result<(), Error> {
+        let input = "3+4";
+        assert_eq!(get_wa_user_short_url(input).await?, "https://is.gd/AZRBWy");
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_wa_user_short_url_6() -> Result<(), Error> {
+        let input = "3 + 4";
+        assert_eq!(get_wa_user_short_url(input).await?, "https://is.gd/lBuhgK");
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_wa_user_short_url_7() -> Result<(), Error> {
+        let input = "test";
+        assert_eq!(get_wa_user_short_url(input).await?, "https://is.gd/NzIyUZ");
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_query_result_parsing() -> Result<(), Error> {
         let body = include_str!("../../tests/resources/wolfram_alpha_api_response.json");
         let _m = mockito::mock("GET", Matcher::Any)
             // Trimmed down version of a full WA response:
@@ -141,4 +238,40 @@ mod tests {
         );
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_query_with_spaces_result_parsing() -> Result<(), Error> {
+        let body = include_str!(
+            "../../tests/resources/wolfram_alpha_api_response_of_input_with_spaces.json"
+        );
+        let _m = mockito::mock("GET", Matcher::Any)
+            // Trimmed down version of a full WA response:
+            .with_body(body)
+            .create();
+        mockito::start();
+
+        let res = wa_query("5 / 10", None, Some(&mockito::server_url())).await?;
+        let res_without_link = res.rsplitn(2, "-").collect::<Vec<&str>>()[1..].join(" ");
+        assert_eq!(
+            res_without_link.trim(),
+            "Exact result: 1/2 - Decimal form: 0.5"
+        );
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_query_with_result_with_primary_pods_parsing() -> Result<(), Error> {
+        let body =
+            include_str!("../../tests/resources/wolfram_alpha_api_response_with_no_primaries.json");
+        let _m = mockito::mock("GET", Matcher::Any)
+            // Trimmed down version of a full WA response:
+            .with_body(body)
+            .create();
+        mockito::start();
+
+        let res = wa_query("what is a url", None, Some(&mockito::server_url())).await?;
+        let res_without_link = res.rsplitn(2, "-").collect::<Vec<&str>>()[1..].join(" ");
+        assert_eq!(res_without_link.trim(), "No primary results.");
+        Ok(())
+    }
 }
diff --git a/src/util/web.rs b/src/util/web.rs
index 77e20a3..bfa7003 100644
--- a/src/util/web.rs
+++ b/src/util/web.rs
@@ -1,5 +1,10 @@
 use anyhow::{Context, Error, Result};
 use reqwest::{get, Url};
+use urlparse::quote_plus as urlparse_quote_plus;
+
+pub(crate) fn quote_plus(text: &str) -> Result<String, Error> {
+    Ok(urlparse_quote_plus(text, b"")?)
+}
 
 // TODO: Either catinator should have a URL shortening utility module,
 // or we should start our own service
@@ -18,3 +23,27 @@ pub(crate) async fn shorten_url(url: &str) -> Result<String, Error> {
 
     Ok(short_url)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{quote_plus};
+    use anyhow::{Error, Result};
+
+    #[test]
+    fn test_quote_plus_1() -> Result<(), Error> {
+        assert_eq!(quote_plus("5/10")?, "5%2F10");
+        Ok(())
+    }
+
+    #[test]
+    fn test_quote_plus_2() -> Result<(), Error> {
+        assert_eq!(quote_plus("1 * 2")?, "1+%2A+2");
+        Ok(())
+    }
+
+    #[test]
+    fn test_quote_plus_3() -> Result<(), Error> {
+        assert_eq!(quote_plus(&quote_plus("1 * 2")?)?, "1%2B%252A%2B2");
+        Ok(())
+    }
+}