From 5e1a98cc9f960d5bb242d52c708cfdfb1c6d3bdb Mon Sep 17 00:00:00 2001
From: Felix Ableitner <me@nutomic.com>
Date: Thu, 23 May 2024 12:56:14 +0200
Subject: [PATCH] more test cases

---
 crates/db_schema/src/newtypes.rs |  2 +-
 crates/federate/src/lib.rs       | 75 +++++++++++++++++++++++++++++---
 2 files changed, 71 insertions(+), 6 deletions(-)

diff --git a/crates/db_schema/src/newtypes.rs b/crates/db_schema/src/newtypes.rs
index 8f2e99cc5..b099c586f 100644
--- a/crates/db_schema/src/newtypes.rs
+++ b/crates/db_schema/src/newtypes.rs
@@ -127,7 +127,7 @@ pub struct LanguageId(pub i32);
 /// The comment reply id.
 pub struct CommentReplyId(i32);
 
-#[derive(Debug, Copy, Clone, Hash, Eq, PartialEq, Serialize, Deserialize, Default)]
+#[derive(Debug, Copy, Clone, Hash, Eq, PartialEq, Serialize, Deserialize, Default, Ord, PartialOrd)]
 #[cfg_attr(feature = "full", derive(DieselNewType, TS))]
 #[cfg_attr(feature = "full", ts(export))]
 /// The instance id.
diff --git a/crates/federate/src/lib.rs b/crates/federate/src/lib.rs
index 87c642998..38c0c235b 100644
--- a/crates/federate/src/lib.rs
+++ b/crates/federate/src/lib.rs
@@ -156,10 +156,16 @@ mod test {
 
   use super::*;
   use activitypub_federation::config::Data;
+  use chrono::DateTime;
+  use lemmy_db_schema::source::{
+    federation_allowlist::FederationAllowList,
+    federation_blocklist::FederationBlockList,
+    instance::InstanceForm,
+  };
   use lemmy_utils::error::LemmyError;
   use serial_test::serial;
   use std::sync::{Arc, Mutex};
-  use tokio::{spawn, time::sleep};
+  use tokio::{spawn, time::sleep};use std::collections::HashSet;
 
   struct TestData {
     send_manager: SendManager,
@@ -213,10 +219,7 @@ mod test {
     }
   }
 
-  // check that correct number of instance workers was started
-  // TODO: need to wrap in Arc or something similar
-  // TODO: test with different `opts`, dead/blocked instances etc
-
+  /// Basic test with default params and only active/allowed instances
   #[tokio::test]
   #[serial]
   async fn test_send_manager() -> LemmyResult<()> {
@@ -224,11 +227,15 @@ mod test {
 
     data.run().await?;
     assert_eq!(3, data.send_manager.workers.len());
+    let workers: HashSet<_> = data.send_manager.workers.keys().cloned().collect();
+    let instances: HashSet<_> = data.instances.iter().map(|i|i.id).collect();
+    assert_eq!(instances, workers);
 
     data.cleanup().await?;
     Ok(())
   }
 
+  /// Running with multiple processes should start correct workers
   #[tokio::test]
   #[serial]
   async fn test_send_manager_processes() -> LemmyResult<()> {
@@ -254,4 +261,62 @@ mod test {
 
     Ok(())
   }
+
+  /// Use blocklist, should not send to blocked instances
+  #[tokio::test]
+  #[serial]
+  async fn test_send_manager_blocked() -> LemmyResult<()> {
+    let mut data = TestData::init(1, 1).await?;
+
+    let domain = data.instances[0].domain.clone();
+    FederationBlockList::replace(&mut data.context.pool(), Some(vec![domain])).await?;
+    data.run().await?;
+    let workers = &data.send_manager.workers;
+    assert_eq!(2, workers.len());
+    assert!(workers.contains_key(&data.instances[1].id));
+    assert!(workers.contains_key(&data.instances[2].id));
+
+    data.cleanup().await?;
+    Ok(())
+  }
+
+  /// Use allowlist, should only send to allowed instance
+  #[tokio::test]
+  #[serial]
+  async fn test_send_manager_allowed() -> LemmyResult<()> {
+    let mut data = TestData::init(1, 1).await?;
+
+    let domain = data.instances[0].domain.clone();
+    FederationAllowList::replace(&mut data.context.pool(), Some(vec![domain])).await?;
+    data.run().await?;
+    let workers = &data.send_manager.workers;
+    assert_eq!(1, workers.len());
+    assert!(workers.contains_key(&data.instances[0].id));
+
+    data.cleanup().await?;
+    Ok(())
+  }
+
+  /// Mark instance as dead, there should be no worker created for it
+  #[tokio::test]
+  #[serial]
+  async fn test_send_manager_dead() -> LemmyResult<()> {
+    let mut data = TestData::init(1, 1).await?;
+
+    let instance = &data.instances[0];
+    let form = InstanceForm::builder()
+      .domain(instance.domain.clone())
+      .updated(DateTime::from_timestamp(0, 0))
+      .build();
+    Instance::update(&mut data.context.pool(), instance.id, form).await?;
+
+    data.run().await?;
+    let workers = &data.send_manager.workers;
+    assert_eq!(2, workers.len());
+    assert!(workers.contains_key(&data.instances[1].id));
+    assert!(workers.contains_key(&data.instances[2].id));
+
+    data.cleanup().await?;
+    Ok(())
+  }
 }