From 80d19073f86dca9f8cff76aa69d5a67bd37687ca Mon Sep 17 00:00:00 2001
From: hi-rustin <rustin.liu@gmail.com>
Date: Wed, 27 Mar 2024 23:21:16 +0800
Subject: [PATCH] test: add an integration test

Signed-off-by: hi-rustin <rustin.liu@gmail.com>
---
 src/typosquat/database.rs              |  4 +--
 src/typosquat/test_util.rs             |  9 +++--
 src/worker/jobs/expiry_notification.rs | 50 ++++++++++++++++++++++++++
 src/worker/jobs/typosquat.rs           |  4 +--
 4 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/src/typosquat/database.rs b/src/typosquat/database.rs
index 62335c0c408..60fdd88a024 100644
--- a/src/typosquat/database.rs
+++ b/src/typosquat/database.rs
@@ -174,8 +174,8 @@ mod tests {
         let mut faker = Faker::new();
 
         // Set up two users.
-        let user_a = faker.user(&mut conn, "a")?;
-        let user_b = faker.user(&mut conn, "b")?;
+        let user_a = faker.user(&mut conn, "a", None)?;
+        let user_b = faker.user(&mut conn, "b", None)?;
 
         // Set up three crates with various ownership schemes.
         let _top_a = faker.crate_and_version(&mut conn, "a", "Hello", &user_a, 2)?;
diff --git a/src/typosquat/test_util.rs b/src/typosquat/test_util.rs
index 04520aca8ff..29cd725f9da 100644
--- a/src/typosquat/test_util.rs
+++ b/src/typosquat/test_util.rs
@@ -101,10 +101,15 @@ impl Faker {
         ))
     }
 
-    pub fn user(&mut self, conn: &mut PgConnection, login: &str) -> anyhow::Result<User> {
+    pub fn user(
+        &mut self,
+        conn: &mut PgConnection,
+        login: &str,
+        email: Option<String>,
+    ) -> anyhow::Result<User> {
         Ok(
             NewUser::new(self.next_id(), login, None, None, "token").create_or_update(
-                None,
+                email.as_deref(),
                 &self.emails,
                 conn,
             )?,
diff --git a/src/worker/jobs/expiry_notification.rs b/src/worker/jobs/expiry_notification.rs
index c5d9f4fd659..4ae946a6790 100644
--- a/src/worker/jobs/expiry_notification.rs
+++ b/src/worker/jobs/expiry_notification.rs
@@ -87,3 +87,53 @@ impl Email for ExpiryNotificationEmail {
         )
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{
+        models::token::ApiToken, schema::api_tokens, test_util::test_db_connection,
+        typosquat::test_util::Faker, util::token::PlainToken,
+    };
+    use diesel::{QueryDsl, SelectableHelper};
+    use lettre::Address;
+
+    #[tokio::test]
+    async fn test_expiry_notification() -> anyhow::Result<()> {
+        let emails = Emails::new_in_memory();
+        let (_test_db, mut conn) = test_db_connection();
+        let mut faker = Faker::new();
+
+        // Set up a user and a token that is about to expire.
+        let user = faker.user(&mut conn, "a", Some("testuser@test.com".to_owned()))?;
+        let token = PlainToken::generate();
+        let expired_at = diesel::dsl::now;
+
+        let token: ApiToken = diesel::insert_into(api_tokens::table)
+            .values((
+                api_tokens::user_id.eq(user.id),
+                api_tokens::name.eq("test_token"),
+                api_tokens::token.eq(token.hashed()),
+                api_tokens::expired_at.eq(expired_at),
+            ))
+            .returning(ApiToken::as_returning())
+            .get_result(&mut conn)?;
+
+        // Check that the token is about to expire.
+        check(&emails, &mut conn)?;
+
+        // Check that an email was sent.
+        let sent_mail = emails.mails_in_memory().unwrap();
+        assert_eq!(sent_mail.len(), 1);
+        let sent = &sent_mail[0];
+        assert_eq!(&sent.0.to(), &["testuser@test.com".parse::<Address>()?]);
+        assert!(sent.1.contains("Your token is about to expire"));
+        let update_token = api_tokens::table
+            .filter(api_tokens::id.eq(token.id))
+            .select(ApiToken::as_select())
+            .first::<ApiToken>(&mut conn)?;
+        assert!(update_token.expiry_notification_at.is_some());
+
+        Ok(())
+    }
+}
diff --git a/src/worker/jobs/typosquat.rs b/src/worker/jobs/typosquat.rs
index f2e5e5b5bca..3d9cbc96c9c 100644
--- a/src/worker/jobs/typosquat.rs
+++ b/src/worker/jobs/typosquat.rs
@@ -130,14 +130,14 @@ mod tests {
         let mut faker = Faker::new();
 
         // Set up a user and a popular crate to match against.
-        let user = faker.user(&mut conn, "a")?;
+        let user = faker.user(&mut conn, "a", None)?;
         faker.crate_and_version(&mut conn, "my-crate", "It's awesome", &user, 100)?;
 
         // Prime the cache so it only includes the crate we just created.
         let cache = Cache::new(vec!["admin@example.com".to_string()], &mut conn)?;
 
         // Now we'll create new crates: one problematic, one not so.
-        let other_user = faker.user(&mut conn, "b")?;
+        let other_user = faker.user(&mut conn, "b", None)?;
         let (angel, _version) = faker.crate_and_version(
             &mut conn,
             "innocent-crate",