From 669a6d84e9f8b348aa5f771392d05e477d5341de Mon Sep 17 00:00:00 2001
From: Yingwen <realevenyag@gmail.com>
Date: Thu, 16 May 2024 19:50:45 +0800
Subject: [PATCH] test: gracefully shutdown postgres client in sql tests
 (#3958)

* chore: debug log

* test: gracefully shutdown pg client
---
 tests-integration/tests/sql.rs | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs
index 97c12f63a5ff..672f2891c5e7 100644
--- a/tests-integration/tests/sql.rs
+++ b/tests-integration/tests/sql.rs
@@ -36,8 +36,12 @@ macro_rules! sql_test {
                         #[$meta]
                     )*
                     async fn [< $test >]() {
+                        common_telemetry::init_default_ut_logging();
+
                         let store_type = tests_integration::test_util::StorageType::$service;
                         if store_type.test_on() {
+                            common_telemetry::info!("test {} starts, store_type: {:?}", stringify!($test), store_type);
+
                             let _ = $crate::sql::$test(store_type).await;
                         }
 
@@ -427,8 +431,10 @@ pub async fn test_postgres_bytea(store_type: StorageType) {
     let (client, connection) = tokio_postgres::connect(&format!("postgres://{addr}/public"), NoTls)
         .await
         .unwrap();
+    let (tx, rx) = tokio::sync::oneshot::channel();
     tokio::spawn(async move {
         connection.await.unwrap();
+        tx.send(()).unwrap();
     });
     let _ = client
         .simple_query("CREATE TABLE test(b BLOB, ts TIMESTAMP TIME INDEX)")
@@ -481,6 +487,9 @@ pub async fn test_postgres_bytea(store_type: StorageType) {
     let val: Vec<u8> = row.get("b");
     assert_eq!(val, [97, 98, 99, 107, 108, 109, 42, 169, 84]);
 
+    drop(client);
+    rx.await.unwrap();
+
     let _ = fe_pg_server.shutdown().await;
     guard.remove_all().await;
 }
@@ -492,8 +501,10 @@ pub async fn test_postgres_datestyle(store_type: StorageType) {
         .await
         .unwrap();
 
+    let (tx, rx) = tokio::sync::oneshot::channel();
     tokio::spawn(async move {
         connection.await.unwrap();
+        tx.send(()).unwrap();
     });
 
     let validate_datestyle = |client: Client, datestyle: &str, is_valid: bool| {
@@ -703,6 +714,9 @@ pub async fn test_postgres_datestyle(store_type: StorageType) {
         }
     }
 
+    drop(client);
+    rx.await.unwrap();
+
     let _ = fe_pg_server.shutdown().await;
     guard.remove_all().await;
 }
@@ -714,8 +728,10 @@ pub async fn test_postgres_timezone(store_type: StorageType) {
         .await
         .unwrap();
 
+    let (tx, rx) = tokio::sync::oneshot::channel();
     tokio::spawn(async move {
         connection.await.unwrap();
+        tx.send(()).unwrap();
     });
 
     let get_row = |mess: Vec<SimpleQueryMessage>| -> String {
@@ -758,6 +774,10 @@ pub async fn test_postgres_timezone(store_type: StorageType) {
             .unwrap(),
     );
     assert_eq!(timezone, "UTC");
+
+    drop(client);
+    rx.await.unwrap();
+
     let _ = fe_pg_server.shutdown().await;
     guard.remove_all().await;
 }
@@ -769,8 +789,10 @@ pub async fn test_postgres_parameter_inference(store_type: StorageType) {
         .await
         .unwrap();
 
+    let (tx, rx) = tokio::sync::oneshot::channel();
     tokio::spawn(async move {
         connection.await.unwrap();
+        tx.send(()).unwrap();
     });
 
     // Create demo table
@@ -796,6 +818,10 @@ pub async fn test_postgres_parameter_inference(store_type: StorageType) {
 
     assert_eq!(1, rows.len());
 
+    // Shutdown the client.
+    drop(client);
+    rx.await.unwrap();
+
     let _ = fe_pg_server.shutdown().await;
     guard.remove_all().await;
 }