From 6bfb8f9e15f2e4d3fddf8208a0ea0776f65fae3e Mon Sep 17 00:00:00 2001
From: Igor Katson <igor.katson@gmail.com>
Date: Wed, 21 Aug 2024 17:39:21 +0100
Subject: [PATCH] E2E test: wait until i am the last task properly

---
 crates/librqbit/src/tests/e2e.rs       | 11 ++++---
 crates/librqbit/src/tests/test_util.rs | 41 +++++++++++++++++++++++++-
 2 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/crates/librqbit/src/tests/e2e.rs b/crates/librqbit/src/tests/e2e.rs
index d762d04..67b6fcb 100644
--- a/crates/librqbit/src/tests/e2e.rs
+++ b/crates/librqbit/src/tests/e2e.rs
@@ -16,7 +16,7 @@ use crate::{
     create_torrent,
     tests::test_util::{
         create_default_random_dir_with_torrents, setup_test_logging, spawn_debug_server,
-        DropChecks, TestPeerMetadata,
+        wait_until_i_am_the_last_task, DropChecks, TestPeerMetadata,
     },
     AddTorrentOptions, AddTorrentResponse, Session, SessionOptions, SessionPersistenceConfig,
 };
@@ -38,10 +38,7 @@ async fn test_e2e_download() {
     .unwrap();
 
     // Wait to ensure everything is dropped.
-    tokio::time::sleep(Duration::from_secs(1)).await;
-
-    let metrics = tokio::runtime::Handle::current().metrics();
-    assert_eq!(metrics.num_alive_tasks(), 1);
+    wait_until_i_am_the_last_task().await;
 
     drop_checks.check().unwrap();
 }
@@ -133,6 +130,9 @@ async fn _test_e2e_download(drop_checks: &DropChecks) {
                     .await
                     .context("error adding torrent")?;
                 let h = handle.into_handle().context("into_handle()")?;
+
+                drop_checks.add(&h.shared, format!("server {i} torrent shared handle"));
+
                 let mut interval = interval(Duration::from_millis(100));
 
                 info!("added torrent");
@@ -145,7 +145,6 @@ async fn _test_e2e_download(drop_checks: &DropChecks) {
                                 if !l.is_finished() {
                                     bail!("torrent went live, but expected it to be finished");
                                 }
-                                drop_checks.add(l, format!("server {i} live"));
                                 Ok(true)
                             }
                             crate::ManagedTorrentState::Error(e) => bail!("error: {e:?}"),
diff --git a/crates/librqbit/src/tests/test_util.rs b/crates/librqbit/src/tests/test_util.rs
index 1c2e79a..051d773 100644
--- a/crates/librqbit/src/tests/test_util.rs
+++ b/crates/librqbit/src/tests/test_util.rs
@@ -2,9 +2,10 @@ use std::{
     io::Write,
     path::Path,
     sync::{Arc, Weak},
+    time::Duration,
 };
 
-use anyhow::Context;
+use anyhow::{bail, Context};
 use axum::{response::IntoResponse, routing::get, Router};
 use librqbit_core::Id20;
 use parking_lot::RwLock;
@@ -165,3 +166,41 @@ impl DropChecks {
         Ok(())
     }
 }
+
+pub async fn wait_until(
+    mut cond: impl FnMut() -> anyhow::Result<()>,
+    timeout: Duration,
+) -> anyhow::Result<()> {
+    let mut interval = tokio::time::interval(Duration::from_millis(10));
+    let mut last_err: Option<anyhow::Error> = None;
+    let res = tokio::time::timeout(timeout, async {
+        loop {
+            interval.tick().await;
+            match cond() {
+                Ok(()) => return Ok::<_, anyhow::Error>(()),
+                Err(e) => last_err = Some(e),
+            }
+        }
+    })
+    .await;
+    if res.is_err() {
+        bail!("wait_until timeout: last result = {last_err:?}")
+    }
+    Ok(())
+}
+
+pub async fn wait_until_i_am_the_last_task() {
+    let metrics = tokio::runtime::Handle::current().metrics();
+    wait_until(
+        || {
+            let num_alive = metrics.num_alive_tasks();
+            if num_alive != 1 {
+                bail!("metrics.num_alive_tasks() = {num_alive}, expected 1")
+            }
+            Ok(())
+        },
+        Duration::from_secs(5),
+    )
+    .await
+    .unwrap();
+}