From 0f2b21491811e269457e9c22fc09a4022e1c0c12 Mon Sep 17 00:00:00 2001
From: "Sijie.Sun" <sunsijie@buaa.edu.cn>
Date: Sat, 24 May 2025 01:22:25 +0800
Subject: [PATCH] fix web test (#872)

---
 easytier-web/src/client_manager/mod.rs     |  7 ++++++
 easytier-web/src/client_manager/session.rs | 27 +++++++++++++++++-----
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/easytier-web/src/client_manager/mod.rs b/easytier-web/src/client_manager/mod.rs
index 2251178f..d8354a34 100644
--- a/easytier-web/src/client_manager/mod.rs
+++ b/easytier-web/src/client_manager/mod.rs
@@ -124,6 +124,7 @@ mod tests {
         },
         web_client::WebClient,
     };
+    use sqlx::Executor;
 
     use crate::{client_manager::ClientManager, db::Db};
 
@@ -133,6 +134,12 @@ mod tests {
         let mut mgr = ClientManager::new(Db::memory_db().await);
         mgr.serve(Box::new(listener)).await.unwrap();
 
+        mgr.db()
+            .inner()
+            .execute("INSERT INTO users (username, password) VALUES ('test', 'test')")
+            .await
+            .unwrap();
+
         let connector = UdpTunnelConnector::new("udp://127.0.0.1:54333".parse().unwrap());
         let _c = WebClient::new(connector, "test", "test");
 
diff --git a/easytier-web/src/client_manager/session.rs b/easytier-web/src/client_manager/session.rs
index 2b26df20..dc59ff95 100644
--- a/easytier-web/src/client_manager/session.rs
+++ b/easytier-web/src/client_manager/session.rs
@@ -69,13 +69,9 @@ struct SessionRpcService {
     data: SharedSessionData,
 }
 
-#[async_trait::async_trait]
-impl WebServerService for SessionRpcService {
-    type Controller = BaseController;
-
-    async fn heartbeat(
+impl SessionRpcService {
+    async fn handle_heartbeat(
         &self,
-        _: BaseController,
         req: HeartbeatRequest,
     ) -> rpc_types::error::Result<HeartbeatResponse> {
         let mut data = self.data.write().await;
@@ -133,6 +129,25 @@ impl WebServerService for SessionRpcService {
     }
 }
 
+#[async_trait::async_trait]
+impl WebServerService for SessionRpcService {
+    type Controller = BaseController;
+
+    async fn heartbeat(
+        &self,
+        _: BaseController,
+        req: HeartbeatRequest,
+    ) -> rpc_types::error::Result<HeartbeatResponse> {
+        let ret = self.handle_heartbeat(req).await;
+        if ret.is_err() {
+            tracing::warn!("Failed to handle heartbeat: {:?}", ret);
+            // sleep for a while to avoid client busy loop
+            tokio::time::sleep(std::time::Duration::from_secs(2)).await;
+        }
+        ret
+    }
+}
+
 pub struct Session {
     rpc_mgr: BidirectRpcManager,