HELP! tokio::test and tokio::main behave differently

[kaixinbaba]

I found the same code, but the performance is different with test run or main run, the following is an example:

async fn mock_fn(addr: &str) -> ZKResult<String> {
    let mut server_list = Vec::new();
    server_list.push(addr.to_string());

    let socket = match TcpStream::connect(server_list.get(0).unwrap().as_str()).await {
        Ok(socket) => socket,
        Err(e) => return Ok(addr.to_string()),
    };

    let (mut reader, mut writer) = io::split(socket);

    tokio::spawn(async move {
        loop {
            info!("writing!!!");
            thread::sleep(Duration::from_millis(500));
        }
        Ok::<_, io::Error>(())
    });
    info!("after spawn");
    Ok(addr.to_string())
}

#[tokio::main]
async fn main() {
    pretty_env_logger::init();
    info!("tokio");
    mock_fn("127.0.0.1:2181").await.unwrap();
    info!("after mock");
}


#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;
    use std::time::Duration;

    #[tokio::test]
    async fn new_test() {
        pretty_env_logger::init();
        info!("tokio");
        mock_fn("127.0.0.1:2181").await.unwrap();
        thread::sleep(Duration::from_secs(10));
        info!("after mock");
    }
}

Here's the output from console by execute fn main:

 INFO  tokio_demo > tokio
 INFO  tokio_demo > after spawn
 INFO  tokio_demo > writing!!!
 INFO  tokio_demo > after mock
 INFO  tokio_demo > writing!!!
 INFO  tokio_demo > writing!!!
 INFO  tokio_demo > writing!!!
 INFO  tokio_demo > writing!!!
 INFO  tokio_demo > writing!!!
 INFO  tokio_demo > writing!!!
...

And here's the output by fn new_test:

 INFO  tokio_demo::tests > tokio
 INFO  tokio_demo        > after spawn
 INFO  tokio_demo::tests > after mock

I have two question:

• You can see that the code in tokio::spawn is not executed at all if fn new_test
• Because I used loop in tokio::spawn, my main function won't stop

So, WHY? Did I miss something?

And now I update the code, add await

async fn mock_fn(addr: &str) -> ZKResult<String> {
    let mut server_list = Vec::new();
    server_list.push(addr.to_string());

    let socket = match TcpStream::connect(server_list.get(0).unwrap().as_str()).await {
        Ok(socket) => socket,
        Err(e) => return Ok(addr.to_string()),
    };

    let (mut reader, mut writer) = io::split(socket);

    tokio::spawn(async move {
        loop {
            info!("writing!!!");
            thread::sleep(Duration::from_millis(500));
        }
        Ok::<_, io::Error>(())
    }).await; // I add await here!!!!
    info!("after spawn");
    Ok(addr.to_string())
}

#[tokio::main]
async fn main() {
    pretty_env_logger::init();
    info!("tokio");
    mock_fn("127.0.0.1:2181").await.unwrap();
    info!("after mock");
}


#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;
    use std::time::Duration;

    #[tokio::test]
    async fn new_test() {
        pretty_env_logger::init();
        info!("tokio");
        mock_fn("127.0.0.1:2181").await.unwrap();
        thread::sleep(Duration::from_secs(10));
        info!("after mock");
    }
}

And now both new_test and main output are same:

 INFO  tokio_demo::tests > tokio (or  INFO  tokio_demo > tokio)
 INFO  tokio_demo        > writing!!!
 INFO  tokio_demo        > writing!!!
 INFO  tokio_demo        > writing!!!
 INFO  tokio_demo        > writing!!!
 INFO  tokio_demo        > writing!!!

BUT!
The code after await will never be executed, because spawn is a loop and will not return, so the code is blocked at await。

My goal is to use a thread loop to process the data from the channel, and the main thread is responsible for sending data to the channel, and the child thread should automatically exit when the main thread ends. How can I implement similar functions in tokio by RUST, I am Rust NOOB, Please help!
I also put this in issue
issue-3337

[Darksonn]

This is because you are blocking the thread and the test runtime is single-threaded. Consider reading Async: What is blocking?. As for the .await on tokio::spawn, it makes mock_fn wait for the spawned task.

[LeoniePhiline]

UPDATE: Answered in Discord. Thanks Alice! The answer is: Must await spawned tasks to ensure they are driven to completion, as the test runtime is shut down once the test returns. In my case, I am adding a dispatch (sync, async) option to my mailer, and dispatching sync in tests.

---

Original question:

I have the case of sending a notification mail on a spawned separate task, while the main task sends a http response. The http response should not wait for the mail to be sent.

This means I am on purpose not .awaiting the tokio::spawn.

In tests, the mail is never sent - the test apparently just succeeds with the main task (http response) finishing.

When testing manually (i.e. sending a real request to the running service), the mail is sent. This means I must have set the test up incorrectly.

I had hoped #[tokio::test(flavor = "multi_thread")] [1] would save me, but that did not help.

How can I make sure all tasks are driven to completion before the test is marked as finished?

I imagine a tokio::try_join-like functionality in the test harness, which might fail on panics or Result::Err on any spawned thread. I am pretty sure this is a common need and I am just overlooking something. Where do I need to look?