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 { 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, 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(()) } }