add heartbeat count

TuGraph-family · Oct 18, 2023 · 6aefbda · 6aefbda
1 parent 58312aa
commit 6aefbda
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 1 deletion.
diff --git a/src/main/java/com/antgroup/tugraph/service/HeartbeatServiceImpl.java b/src/main/java/com/antgroup/tugraph/service/HeartbeatServiceImpl.java
@@ -13,10 +13,13 @@ public TuGraphDBManagement.HeartbeatResponse detectHeartbeat(TuGraphDBManagement
         String reqMsg = request.getRequestMsg();
         log.info(reqMsg);
         String respMsg = "This is a heartbeat response message.";
+        int heartbeatCount = request.getHeartbeatCount();
         TuGraphDBManagement.HeartbeatResponse.Builder respBuilder =
             TuGraphDBManagement.HeartbeatResponse
                 .newBuilder();
-        respBuilder.setResponseMsg(respMsg);
+        respBuilder
+            .setResponseMsg(respMsg)
+            .setHeartbeatCount(heartbeatCount + 1);
         return respBuilder.build();
     }
 

diff --git a/src/main/proto/tugraph_db_management.proto b/src/main/proto/tugraph_db_management.proto
@@ -120,10 +120,12 @@ message JobManagementResponse {
 
 message HeartbeatRequest {
       required string request_msg = 1;
+      required int32 heartbeat_count = 2;
 };
 
 message HeartbeatResponse {
       required string response_msg = 1;
+      required int32 heartbeat_count = 2;
 };
 
 service JobManagementService {

diff --git a/src/test/java/com/antgroup/tugraph/TuGraphDBManagementApplicationTests.java b/src/test/java/com/antgroup/tugraph/TuGraphDBManagementApplicationTests.java
@@ -23,6 +23,9 @@ class TuGraphDBManagementApplicationTests {
 	@Autowired
 	private JobManagementServiceImpl jobManagementService;
 
+	@Autowired
+	private HeartbeatServiceImpl heartbeatService;
+
 	@Autowired
     private JobService jobService;
 
@@ -167,4 +170,25 @@ void testJobManagementService() {
 		assertEquals(0, JobList.size());
 	}
 
+	@Test
+	@Transactional
+	@Order(1)
+	void testHeartbeatService() {
+		log.info("start testing heartbeat service.");
+
+		String reqMsg = "This is a heartbeat request message.";
+        int heartbeatCount = 0;
+
+		TuGraphDBManagement.HeartbeatRequest heartbeatRequest =
+            TuGraphDBManagement.HeartbeatRequest
+                .newBuilder()
+				.setRequestMsg(reqMsg)
+				.setHeartbeatCount(heartbeatCount)
+				.build();
+        TuGraphDBManagement.HeartbeatResponse heartbeatResponse =
+			heartbeatService.detectHeartbeat(heartbeatRequest);
+
+		assertEquals(heartbeatCount + 1, heartbeatResponse.getHeartbeatCount());
+	}
+
 }