add tests

postgresml · chillenberger · Jun 14, 2024 · Jun 16, 2024 · Jun 16, 2024 · Jun 17, 2024
commit e1f64d66863ca8c65b4b82f1a2df04095902f9d3
diff --git a/pgml-dashboard/src/lib.rs b/pgml-dashboard/src/lib.rs
@@ -390,7 +390,6 @@ pub fn replace_banner_product(
     context: &Cluster,
 ) -> Result<Response, Error> {
     let mut all_notification_cookies = Notifications::get_viewed(cookies);
-
     let current_notification_cookie = all_notification_cookies.iter().position(|x| x.id == id);
 
     match current_notification_cookie {
@@ -408,8 +407,8 @@ pub fn replace_banner_product(
 
     Notifications::update_viewed(&all_notification_cookies, cookies);
 
-    // Get the notification that triggered this call.
-    // Guaranteed to exist since it built the component that called this, so this is safe to unwrap.
+    // Get the notification that triggered this call..
+    // unwrap notifications if fine since we should panic if this is missing.
     let last_notification = context
         .notifications
         .as_ref()
@@ -424,11 +423,15 @@ pub fn replace_banner_product(
             .into_iter()
             .filter(|n: &Notification| -> bool {
                 let n = n.clone().set_viewed(n.id == id);
-                Notification::product_filter(
-                    &n,
-                    last_notification.clone().unwrap().level.clone(),
-                    deployment_id.clone(),
-                )
+                if last_notification.clone().is_none() {
+                    return false;
+                } else {
+                    Notification::product_filter(
+                        &n,
+                        last_notification.clone().unwrap().level.clone(),
+                        deployment_id.clone(),
+                    )
+                }
             })
             .next(),
         _ => None,
@@ -519,3 +522,176 @@ pub fn routes() -> Vec<Route> {
 pub async fn migrate(pool: &PgPool) -> anyhow::Result<()> {
     Ok(sqlx::migrate!("./migrations").run(pool).await?)
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::components::sections::footers::MarketingFooter;
+    use crate::guards::Cluster;
+    use crate::utils::config;
+    use rocket::fairing::AdHoc;
+    use rocket::http::{Cookie, Status};
+    use rocket::local::asynchronous::Client;
+
+    #[sqlx::test]
+    async fn test_remove_modal() {
+        let rocket = rocket::build().mount("/", routes());
+        let client = Client::untracked(rocket).await.unwrap();
+
+        let cookie = vec![
+            NotificationCookie {
+                id: "1".to_string(),
+                time_viewed: Some(chrono::Utc::now() - chrono::Duration::days(1)),
+                time_modal_viewed: Some(chrono::Utc::now() - chrono::Duration::days(1)),
+            },
+            NotificationCookie {
+                id: "2".to_string(),
+                time_viewed: None,
+                time_modal_viewed: None,
+            },
+        ];
+
+        let response = client
+            .get("/notifications/product/modal/remove_modal?id=1")
+            .private_cookie(Cookie::new("session", Notifications::safe_serialize_session(&cookie)))
+            .dispatch()
+            .await;
+
+        let time_modal_viewed = Notifications::get_viewed(response.cookies())
+            .get(0)
+            .unwrap()
+            .time_modal_viewed;
+
+        // Update modal view time for existing notification cookie
+        assert_eq!(time_modal_viewed.is_some(), true);
+
+        let response = client
+            .get("/notifications/product/modal/remove_modal?id=3")
+            .private_cookie(Cookie::new("session", Notifications::safe_serialize_session(&cookie)))
+            .dispatch()
+            .await;
+
+        let time_modal_viewed = Notifications::get_viewed(response.cookies())
+            .get(0)
+            .unwrap()
+            .time_modal_viewed;
+
+        // Update modal view time for new notification cookie
+        assert_eq!(time_modal_viewed.is_some(), true);
+    }
+
+    #[sqlx::test]
+    async fn test_remove_banner_product() {
+        let rocket = rocket::build().mount("/", routes());
+        let client = Client::untracked(rocket).await.unwrap();
+
+        let cookie = vec![
+            NotificationCookie {
+                id: "1".to_string(),
+                time_viewed: Some(chrono::Utc::now() - chrono::Duration::days(1)),
+                time_modal_viewed: Some(chrono::Utc::now() - chrono::Duration::days(1)),
+            },
+            NotificationCookie {
+                id: "2".to_string(),
+                time_viewed: None,
+                time_modal_viewed: Some(chrono::Utc::now() - chrono::Duration::days(1)),
+            },
+        ];
+
+        let response = client
+            .get("/notifications/product/remove_banner?id=1&target=ajskghjfbs")
+            .private_cookie(Cookie::new("session", Notifications::safe_serialize_session(&cookie)))
+            .dispatch()
+            .await;
+
+        let time_viewed = Notifications::get_viewed(response.cookies())
+            .get(0)
+            .unwrap()
+            .time_viewed;
+
+        // Update view time for existing notification cookie
+        assert_eq!(time_viewed.is_some(), true);
+
+        let response = client
+            .get("/notifications/product/remove_banner?id=3&target=ajfadghs")
+            .private_cookie(Cookie::new("session", Notifications::safe_serialize_session(&cookie)))
+            .dispatch()
+            .await;
+
+        let time_viewed = Notifications::get_viewed(response.cookies())
+            .get(0)
+            .unwrap()
+            .time_viewed;
+
+        // Update view time for new notification cookie
+        assert_eq!(time_viewed.is_some(), true);
+    }
+
+    #[sqlx::test]
+    async fn test_replace_banner_product() {
+        let notification1 = Notification::new("Test notification 1")
+            .set_level(&NotificationLevel::ProductMedium)
+            .set_deployment("1");
+        let notification2 = Notification::new("Test notification 2")
+            .set_level(&NotificationLevel::ProductMedium)
+            .set_deployment("1");
+        let notification3 = Notification::new("Test notification 3").set_level(&NotificationLevel::ProductMarketing);
+
+        let rocket = rocket::build()
+            .attach(AdHoc::on_request("request", |req, _| {
+                Box::pin(async {
+                    req.local_cache(|| Cluster {
+                        pool: None,
+                        context: Context {
+                            user: models::User::default(),
+                            cluster: models::Cluster::default(),
+                            dropdown_nav: StaticNav { links: vec![] },
+                            product_left_nav: StaticNav { links: vec![] },
+                            marketing_footer: MarketingFooter::new().render_once().unwrap(),
+                            head_items: None,
+                        },
+                        notifications: Some(vec![
+                            Notification::new("Test notification 1")
+                                .set_level(&NotificationLevel::ProductMedium)
+                                .set_deployment("1"),
+                            Notification::new("Test notification 2")
+                                .set_level(&NotificationLevel::ProductMedium)
+                                .set_deployment("1"),
+                            Notification::new("Test notification 3").set_level(&NotificationLevel::ProductMarketing),
+                        ]),
+                    });
+                })
+            }))
+            .mount("/", routes());
+
+        let client = Client::tracked(rocket).await.unwrap();
+
+        let response = client
+            .get(format!(
+                "/notifications/product/replace_banner?id={}&deployment_id=1",
+                notification1.id
+            ))
+            .dispatch()
+            .await;
+
+        let body = response.into_string().await.unwrap();
+        let rsp_contains_next_notification = body.contains("Test notification 2");
+
+        // ensure banner is replaced with next notification of same type
+        assert_eq!(rsp_contains_next_notification, true);
+
+        let response = client
+            .get(format!(
+                "/notifications/product/replace_banner?id={}&deployment_id=1",
+                notification2.id
+            ))
+            .dispatch()
+            .await;
+
+        let body = response.into_string().await.unwrap();
+        let rsp_contains_next_notification = body.contains("Test notification 3");
+
+        // ensure next notification is not found
+        assert_eq!(rsp_contains_next_notification, false);
+    }
+}
diff --git a/pgml-dashboard/src/utils/cookies.rs b/pgml-dashboard/src/utils/cookies.rs
@@ -2,7 +2,7 @@ use chrono;
 use rocket::http::{Cookie, CookieJar};
 use rocket::serde::{Deserialize, Serialize};
 
-#[derive(Debug, Clone, Deserialize, Serialize)]
+#[derive(Debug, Clone, Deserialize, Serialize, Default, PartialEq)]
 pub struct NotificationCookie {
     pub id: String,
     pub time_viewed: Option<chrono::DateTime<chrono::Utc>>,
@@ -12,51 +12,113 @@ pub struct NotificationCookie {
 pub struct Notifications {}
 
 impl Notifications {
-    pub fn update_viewed(new: &Vec<NotificationCookie>, cookies: &CookieJar<'_>) {
-        let serialized = new
-            .iter()
-            .map(|x| serde_json::to_string(x).unwrap())
-            .collect::<Vec<String>>();
+    pub fn update_viewed(all_desired_notifications: &Vec<NotificationCookie>, cookies: &CookieJar<'_>) {
+        let session = Notifications::safe_serialize_session(all_desired_notifications);
 
-        let mut cookie = Cookie::new("session", format!(r#"{{"notifications": [{}]}}"#, serialized.join(",")));
+        let mut cookie = Cookie::new("session", session);
         cookie.set_max_age(::time::Duration::weeks(4));
         cookies.add_private(cookie);
     }
 
     pub fn get_viewed(cookies: &CookieJar<'_>) -> Vec<NotificationCookie> {
-        let viewed: Vec<NotificationCookie> = match cookies.get_private("session") {
-            Some(session) => {
-                match serde_json::from_str::<serde_json::Value>(session.value())
-                    .unwrap_or_else(|_| serde_json::from_str::<serde_json::Value>(r#"{"notifications": []}"#).unwrap())
-                    ["notifications"]
-                    .as_array()
-                {
-                    Some(items) => items
-                        .into_iter()
-                        .map(|x| {
-                            serde_json::from_str::<NotificationCookie>(&x.to_string()).unwrap_or_else(|_| {
-                                serde_json::from_str::<String>(&x.to_string())
-                                    .and_then(|z| {
-                                        Ok(NotificationCookie {
-                                            id: z,
-                                            time_viewed: None,
-                                            time_modal_viewed: None,
-                                        })
-                                    })
-                                    .unwrap_or_else(|_| NotificationCookie {
-                                        id: "".to_string(),
-                                        time_viewed: None,
-                                        time_modal_viewed: None,
-                                    })
-                            })
-                        })
-                        .collect::<Vec<NotificationCookie>>(),
-                    _ => vec![],
-                }
-            }
+        match cookies.get_private("session") {
+            Some(session) => Notifications::safe_deserialize_session(session.value()),
             None => vec![],
-        };
+        }
+    }
+
+    pub fn safe_deserialize_session(session: &str) -> Vec<NotificationCookie> {
+        match serde_json::from_str::<serde_json::Value>(session).unwrap_or_else(|_| {
+            serde_json::from_str::<serde_json::Value>(&Notifications::safe_serialize_session(&vec![])).unwrap()
+        })["notifications"]
+            .as_array()
+        {
+            Some(items) => items
+                .into_iter()
+                .map(|notification| {
+                    serde_json::from_str::<NotificationCookie>(&notification.to_string()).unwrap_or_else(|_| {
+                        serde_json::from_str::<String>(&notification.to_string())
+                            .and_then(|id| {
+                                Ok(NotificationCookie {
+                                    id,
+                                    time_viewed: None,
+                                    time_modal_viewed: None,
+                                })
+                            })
+                            .unwrap_or_else(|_| NotificationCookie::default())
+                    })
+                })
+                .collect::<Vec<NotificationCookie>>(),
+            _ => vec![],
+        }
+    }
+
+    pub fn safe_serialize_session(cookies: &Vec<NotificationCookie>) -> String {
+        let serialized = cookies
+            .iter()
+            .map(|x| serde_json::to_string(x))
+            .filter(|x| x.is_ok())
+            .map(|x| x.unwrap())
+            .collect::<Vec<String>>();
+
+        format!(r#"{{"notifications": [{}]}}"#, serialized.join(","))
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    // Test that we can safely deserialize expected session data.
+    #[test]
+    fn test_safe_deserialize_session() {
+        let session = r#"{"notifications": [{"id": "1", "time_viewed": null, "time_modal_viewed": null}, {"id": "1234567891234", "time_viewed": "2021-08-01T00:00:00Z"}]}"#;
+        let expected = vec![
+            NotificationCookie {
+                id: "1".to_string(),
+                time_viewed: None,
+                time_modal_viewed: None,
+            },
+            NotificationCookie {
+                id: "1234567891234".to_string(),
+                time_viewed: Some(
+                    chrono::DateTime::parse_from_rfc3339("2021-08-01T00:00:00Z")
+                        .unwrap()
+                        .into(),
+                ),
+                time_modal_viewed: None,
+            },
+        ];
+        assert_eq!(Notifications::safe_deserialize_session(session), expected);
+    }
+
+    // Test that new notification system is backwards compatible.
+    #[test]
+    fn test_safe_deserialize_session_old_form() {
+        let session = r#"{"notifications": ["123456789"]}"#;
+        let expected = vec![NotificationCookie {
+            id: "123456789".to_string(),
+            time_viewed: None,
+            time_modal_viewed: None,
+        }];
+        assert_eq!(Notifications::safe_deserialize_session(session), expected);
+    }
+
+    #[test]
+    fn test_safe_deserialize_session_empty() {
+        let session = r#"{}"#;
+        let expected: Vec<NotificationCookie> = vec![];
+        assert_eq!(Notifications::safe_deserialize_session(session), expected);
+    }
 
-        viewed
+    #[test]
+    fn test_safe_serialize_session() {
+        let cookies = vec![NotificationCookie {
+            id: "1".to_string(),
+            time_viewed: None,
+            time_modal_viewed: None,
+        }];
+        let expected = r#"{"notifications": [{"id":"1","time_viewed":null,"time_modal_viewed":null}]}"#;
+        assert_eq!(Notifications::safe_serialize_session(&cookies), expected);
     }
 }