From ec4402df419624971dc62fc70b8349a3bb3f4ab6 Mon Sep 17 00:00:00 2001
From: lklimek <842586+lklimek@users.noreply.github.com>
Date: Mon, 1 Jul 2024 17:14:41 +0200
Subject: [PATCH] test(abci): fix tokio spawning in grpc test (#81)

* test(abci): fix tokio spawning in grpc test

* chore: minor refactor
---
 abci/Cargo.toml             |  1 +
 abci/tests/common/docker.rs | 32 +++++++++++++++++++++-----------
 abci/tests/grpc.rs          | 13 +++++--------
 3 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/abci/Cargo.toml b/abci/Cargo.toml
index 2be56cc..890056a 100644
--- a/abci/Cargo.toml
+++ b/abci/Cargo.toml
@@ -87,3 +87,4 @@ hex = { version = "0.4.3" }
 lazy_static = { version = "1.4.0" }
 # For tests of gRPC server
 tonic = { version = "0.11.0" }
+pollster = { version = "0.3.0" }
diff --git a/abci/tests/common/docker.rs b/abci/tests/common/docker.rs
index 1645dc4..f8b9c88 100644
--- a/abci/tests/common/docker.rs
+++ b/abci/tests/common/docker.rs
@@ -95,7 +95,7 @@ impl TenderdashDocker {
             tokio::signal::ctrl_c().await.unwrap();
             error!("Received ctrl+c, removing Tenderdash container");
 
-            let stopped = timeout(Duration::from_secs(15), Self::stop(id, &docker))
+            let stopped = timeout(Duration::from_secs(15), Self::stop(id, docker))
                 .await
                 .expect("timeout removing tenderdash container");
             if stopped.is_err() {
@@ -210,24 +210,28 @@ impl TenderdashDocker {
     /// Print 200 most recent logs from Tenderdash on standard error.
     #[allow(dead_code)]
     pub fn print_logs(&self) {
-        let id = &self.id;
+        let id = self.id.clone();
 
         if !id.is_empty() {
             debug!("Printing Tenderdash logs");
             let rt = &self.runtime;
-            let docker = &self.docker;
+            let docker = self.docker.clone();
 
-            rt.block_on(Self::emit_logs(id, docker))
+            // We use `pollster` to block on the future, as tokio might not block_on()
+            // when called inside a running runtime.
+            let fut = rt.spawn(Self::emit_logs(id, docker));
+            pollster::block_on(fut)
+                .expect("cannot spawn log emitting fn")
                 .expect("cannot emit logs");
         }
     }
 
-    async fn emit_logs(id: &str, docker: &Docker) -> Result<(), anyhow::Error> {
+    async fn emit_logs(id: String, docker: Docker) -> Result<(), anyhow::Error> {
         let stderror = tokio::io::stderr();
         let mut dest = tokio::io::BufWriter::new(stderror);
 
         let mut logs = docker.logs(
-            id,
+            &id,
             Some(bollard::container::LogsOptions {
                 follow: false,
                 stdout: true,
@@ -251,7 +255,7 @@ impl TenderdashDocker {
         Ok(())
     }
 
-    async fn stop(id: String, docker: &Docker) -> Result<(), anyhow::Error> {
+    async fn stop(id: String, docker: Docker) -> Result<(), anyhow::Error> {
         debug!("Stopping Tenderdash container");
         docker
             .remove_container(
@@ -270,9 +274,13 @@ impl TenderdashDocker {
 impl Drop for TenderdashDocker {
     fn drop(&mut self) {
         if !self.id.is_empty() {
-            let _ = self
-                .runtime
-                .block_on(Self::stop(self.id.clone(), &self.docker));
+            let id = self.id.clone();
+            let docker = self.docker.clone();
+
+            let fut = self.runtime.spawn(Self::stop(id, docker));
+            pollster::block_on(fut)
+                .expect("cannot stop container")
+                .expect("cannot stop container");
         }
     }
 }
@@ -281,9 +289,11 @@ impl Drop for TenderdashDocker {
 #[allow(dead_code)]
 pub fn setup_td_logs_panic(td_docker: &Arc<TenderdashDocker>) {
     let weak_ref = Arc::downgrade(td_docker);
-    std::panic::set_hook(Box::new(move |_| {
+    let original = std::panic::take_hook();
+    std::panic::set_hook(Box::new(move |info| {
         if let Some(td) = weak_ref.upgrade() {
             td.print_logs()
         }
+        original(info);
     }));
 }
diff --git a/abci/tests/grpc.rs b/abci/tests/grpc.rs
index cf600be..6754a0f 100644
--- a/abci/tests/grpc.rs
+++ b/abci/tests/grpc.rs
@@ -21,7 +21,7 @@ use tenderdash_abci::proto;
 use tonic::{async_trait, Response, Status};
 
 #[cfg(feature = "docker-tests")]
-#[tokio::test]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 /// Test server listening on ipv4 address.
 ///
 /// See [tcp_server_test()].
@@ -34,7 +34,7 @@ async fn test_ipv4_server() {
 }
 
 #[cfg(feature = "docker-tests")]
-#[tokio::test]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[ignore = "IPv6 does not work for gRPC, most likely bug on Tenderdash side"]
 /// Test server listening on ipv6 address.
 ///
@@ -74,7 +74,7 @@ async fn grpc_server_test(test_name: &str, bind_address: &str) {
     let addr = bind_address.parse().expect("address must be valid");
     let server_cancel = cancel.clone();
     let server_handle = tokio::spawn(async move {
-        tracing::debug!("starting gRPC server");
+        tracing::debug!(?addr, "starting gRPC server");
         Server::builder()
             .add_service(AbciApplicationServer::new(app))
             .serve_with_shutdown(addr, server_cancel.cancelled())
@@ -87,7 +87,7 @@ async fn grpc_server_test(test_name: &str, bind_address: &str) {
     let container_name = format!("tenderdash_{}", test_name);
 
     let td = tokio::task::spawn_blocking(move || {
-        tracing::debug!("starting Tenderdash in Docker container");
+        tracing::debug!(addr=?socket_uri, "starting Tenderdash in Docker container");
         let td = Arc::new(common::docker::TenderdashDocker::new(
             &container_name,
             None,
@@ -113,10 +113,7 @@ async fn grpc_server_test(test_name: &str, bind_address: &str) {
         }
     }
 
-    tokio::task::spawn_blocking(move || drop(td))
-        .await
-        .expect("tenderdash cleanup");
-
+    drop(td);
     tracing::info!("Test finished successfully");
 }