Issue-185 - Fix hanging tests

Author: huntharo

extension/tests/lambda_service.rs

@@ -796,15 +796,15 @@ async fn test_lambda_service_loop_100_valid_get_requests() {
   });
 
   // Let the router run wild
-  tokio::spawn(async move {
+  for _ in 0..100 {
     mock_router_server
       .release_request_tx
       .lock()
       .await
       .send(())
       .await
       .unwrap();
-  });
+  }
 
   let mut options = Options::default();
 
@@ -882,19 +882,18 @@ async fn test_lambda_service_loop_100_valid_post_requests() {
     .mount(&mock_app_server)
     .await;
 
-  // Let the router run wild
-  tokio::spawn(async move {
+  // Push 100 messages into the channel (it's limit)
+  for _ in 0..100 {
     mock_router_server
       .release_request_tx
       .lock()
       .await
       .send(())
       .await
       .unwrap();
-  });
+  }
 
   let mut options = Options::default();
-
   options.port = mock_app_server.address().port();
   let initialized = true;
 
@@ -1549,16 +1548,15 @@ async fn test_lambda_service_loop_100_requests_contained_app_connection_close_he
       .body("Bananas");
   });
 
-  // Let the router run wild
-  tokio::spawn(async move {
+  for _ in 0..100 {
     mock_router_server
       .release_request_tx
       .lock()
       .await
       .send(())
       .await
       .unwrap();
-  });
+  }
 
   let mut options = Options::default();
 
@@ -2117,7 +2115,6 @@ async fn test_lambda_service_router_connects_app_crashes_graceful_close() {
 
 #[tokio::test]
 async fn test_lambda_service_router_closes_quickly_when_no_requests() {
-  env_logger::init();
   // Start router server
   let mock_router_server =
     mock_router::setup_router(mock_router::RouterParamsBuilder::new().build());

extension/tests/router_channel.rs

@@ -456,17 +456,17 @@ async fn test_channel_invalid_request_headers_should_continue() {
   });
   let app_endpoint: Endpoint = mock_app_server.base_url().parse().unwrap();
 
-  // Release the request after a few seconds
-  tokio::spawn(async move {
-    tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
+  for _ in 0..2 {
+    // Release the request after a few seconds
+    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
     mock_router_server
       .release_request_tx
       .lock()
       .await
       .send(())
       .await
       .unwrap();
-  });
+  }
 
   let app_client = create_app_client();
   let router_client = create_router_client();

extension/tests/support/mock_router.rs

@@ -313,8 +313,6 @@ pub fn setup_router(params: RouterParams) -> RouterResult {
 
                             // Send static request to extension
                             if request_method == RequestMethod::ShutdownWithoutResponse {
-                              release_request_rx.lock().await.recv().await;
-
                               // Close the body stream
                               tx.shutdown().await.unwrap();
                               return;
@@ -343,7 +341,7 @@ pub fn setup_router(params: RouterParams) -> RouterResult {
                                 let data = b"GET /bananas_query_unencoded_brackets?cat=[dog]&cat=log&cat=cat HTTP/1.1\r\nHost: localhost\r\nTest-Header: foo\r\n\r\n";
                                 tx.write_all(data).await.unwrap();
                             } else if request_method == RequestMethod::GetGoAwayOnBody {
-                                println!("Sending GoAway on body");
+                                println!("{} Sending GoAway on body, request_count: {}", chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"), request_count);
                                 let data = b"GET /_lambda_dispatch/goaway HTTP/1.1\r\nHost: localhost\r\nTest-Header: foo\r\n\r\n";
                                 tx.write_all(data).await.unwrap();
                             } else if request_method == RequestMethod::GetInvalidHeaders {