From 86d3550901c2f2eadc70e47fa811a588c87da4a6 Mon Sep 17 00:00:00 2001
From: Felix Ableitner <me@nutomic.com>
Date: Thu, 23 May 2024 11:35:21 +0200
Subject: [PATCH] basic test working

---
 crates/db_schema/src/impls/instance.rs |  6 +++++-
 crates/federate/src/lib.rs             | 25 ++++++++++++++-----------
 2 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/crates/db_schema/src/impls/instance.rs b/crates/db_schema/src/impls/instance.rs
index 67da14823..f5dcdd9ef 100644
--- a/crates/db_schema/src/impls/instance.rs
+++ b/crates/db_schema/src/impls/instance.rs
@@ -94,11 +94,15 @@ impl Instance {
       .await
   }
 
-  #[cfg(test)]
+  /// Only for use in tests
   pub async fn delete_all(pool: &mut DbPool<'_>) -> Result<usize, Error> {
     let conn = &mut get_conn(pool).await?;
+    diesel::delete(federation_queue_state::table)
+      .execute(conn)
+      .await?;
     diesel::delete(instance::table).execute(conn).await
   }
+
   pub async fn allowlist(pool: &mut DbPool<'_>) -> Result<Vec<Self>, Error> {
     let conn = &mut get_conn(pool).await?;
     instance::table
diff --git a/crates/federate/src/lib.rs b/crates/federate/src/lib.rs
index 9c2fdf57d..b279701ba 100644
--- a/crates/federate/src/lib.rs
+++ b/crates/federate/src/lib.rs
@@ -58,11 +58,11 @@ impl SendManager {
   }
 
   pub fn run(mut self) -> CancellableTask {
-    let task = CancellableTask::spawn(WORKER_EXIT_TIMEOUT, move |cancel| async move {
+    let cancel = CancellableTask::spawn(WORKER_EXIT_TIMEOUT, move |cancel| async move {
       self.do_loop(cancel).await.unwrap();
       self.cancel().await.unwrap();
     });
-    task
+    cancel
   }
 
   async fn do_loop(&mut self, cancel: CancellationToken) -> LemmyResult<()> {
@@ -152,7 +152,7 @@ impl SendManager {
 mod test {
 
   use super::*;
-  use tokio::time::sleep;
+  use tokio::{spawn, time::sleep};
 
   #[tokio::test]
   async fn test_start_stop_federation_workers() -> LemmyResult<()> {
@@ -176,20 +176,23 @@ mod test {
     ];
 
     // start it and wait a moment
-    let task = SendManager::new(opts, federation_config);
-    task.run();
-    sleep(Duration::from_secs(1));
+    let mut task = SendManager::new(opts, federation_config);
+    let cancel = CancellationToken::new();
+    let cancel_ = cancel.clone();
+    spawn(async move {
+      sleep(Duration::from_millis(100)).await;
+      cancel_.cancel();
+    });
+    task.do_loop(cancel.clone()).await.unwrap();
+    assert_eq!(3, task.workers.len());
 
     // 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
-    assert_eq!(3, task.workers.len());
+    //assert_eq!(3, task.workers.len());
 
     // cleanup
-    for i in instances {
-      Instance::delete(pool, i.id).await?;
-    }
-    task.cancel().await?;
+    Instance::delete_all(pool).await?;
     Ok(())
   }
 }