Merge pull request #1551 from st1020/fix/heartbeat_handler_save_fail

fix: fix heartbeat handler save fail
HXSecurity · Jun 26, 2023 · c14d47d · c14d47d
2 parents 9de217d + 2380317
commit c14d47d
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 19 deletions.
diff --git a/dongtai_protocol/report/handler/heartbeat_handler.py b/dongtai_protocol/report/handler/heartbeat_handler.py
@@ -13,13 +13,11 @@
 from dongtai_common.models.vulnerablity import IastVulnerabilityModel
 from dongtai_common.utils import const
 from django.utils.translation import gettext_lazy as _
-from dongtai_common.models.server import IastServer
 from dongtai_protocol.report.handler.report_handler_interface import IReportHandler
 from dongtai_protocol.report.report_handler_factory import ReportHandler
-from django.db.models import (QuerySet, Q, F)
-from dongtai_common.models.project import IastProject, VulValidation
+from django.db.models import (QuerySet, Q)
+from dongtai_common.models.project import VulValidation
 from dongtai_common.utils.systemsettings import get_vul_validate
-from dongtai_common.models.agent import IastAgent
 from django.core.cache import cache
 
 logger = logging.getLogger('dongtai.openapi')
@@ -77,14 +75,14 @@ def save_heartbeat(self):
             default_dict['report_queue'] = self.report_queue
             default_dict['method_queue'] = self.method_queue
             default_dict['replay_queue'] = self.replay_queue
-            IastHeartbeat.objects.update_or_create(agent_id=self.agent_id,
-                                                   defaults=default_dict)
         elif self.return_queue == 0:
             if self.req_count is not None:
                 default_dict['req_count'] = self.req_count
             default_dict['memory'] = self.memory
             default_dict['cpu'] = self.cpu
             default_dict['disk'] = self.disk
+            IastHeartbeat.objects.update_or_create(agent_id=self.agent_id,
+                                                   defaults=default_dict)
         else:
             default_dict['memory'] = self.memory
             default_dict['cpu'] = self.cpu

diff --git a/test/apiserver/test_agent_heartbeat.py b/test/apiserver/test_agent_heartbeat.py
@@ -1,27 +1,36 @@
 from test.apiserver.test_agent_base import AgentTestCase
 from dongtai_engine.tasks import is_alive
-from dongtai_common.models.agent import IastAgent
 from dongtai_common.models.heartbeat import IastHeartbeat
 from time import time
 
 
 class ApiHeartBeatTestCase(AgentTestCase):
-
     def test_agent_heart_beat(self):
         data = {
-            "detail": {
+            "detail": {
                 #                'isCoreInstalled': 1,
-                'agentId': 28,
-                'disk': '{"rate":71}',
-                'memory': '{"total":"875MB","rate":7,"use":"65.872MB"}',
-                'performance':
-                '[{"metricsKey":"cpuUsage","collectDate":"2022-12-14 12:02:05.009","metricsValue":{"cpuUsagePercentage":0.16666666666667052}},{"metricsKey":"memoryUsage","collectDate":"2022-12-14 12:02:05.009","metricsValue":{"init":65011712,"systemMaxLimit":9223372036854771712,"committed":124256256,"memUsagePercentage":9.841756766183035,"max":917504000,"used":90298512}},{"metricsKey":"memoryNoHeapUsage","collectDate":"2022-12-14 12:02:05.009","metricsValue":{"init":4121784320,"systemMaxLimit":9223372036854771712,"committed":3886866432,"memUsagePercentage":94.30057786235646,"max":4121784320,"used":3886866432}},{"metricsKey":"garbageInfo","collectDate":"2022-12-14 12:02:05.009","metricsValue":{"collectionInfoList":[{"collectionTime":18049,"collectionCount":2715,"tenured":false,"collectionName":"PS Scavenge"},{"collectionTime":948,"collectionCount":4,"tenured":true,"collectionName":"PS MarkSweep"}]}},{"metricsKey":"threadInfo","collectDate":"2022-12-14 12:02:05.225","metricsValue":{"dongTaiThreadInfoList":[{"cpuUsage":0,"cpuTime":29161574907,"name":"DongTai-IAST-ServerConfigMonitor","id":14},{"cpuUsage":0,"cpuTime":27679480672,"name":"DongTai-IAST-ConfigMonitor","id":15},{"cpuUsage":0.035035544200697016,"cpuTime":13643355644,"name":"DongTai-IAST-PerformanceMonitor","id":16},{"cpuUsage":0,"cpuTime":111735614230,"name":"DongTai-IAST-EngineMonitor","id":17},{"cpuUsage":0,"cpuTime":61038194468,"name":"DongTai-IAST-HearBeatMonitor","id":18},{"cpuUsage":0,"cpuTime":303498408,"name":"DongTai-IAST-SecondFallbackMonitor","id":19},{"cpuUsage":0,"cpuTime":859911477,"name":"DongTai-IAST-DongTaiThreadMonitor","id":20},{"cpuUsage":0,"cpuTime":26344322815,"name":"DongTai-IAST-HeartBeat","id":64}],"threadCount":51,"dongTaiThreadCount":8,"daemonThreadCount":47,"peakThreadCount":53}}]',
+                "agentId": 28,
+                "disk": '{"rate":71}',
+                "memory": '{"total":"875MB","rate":7,"use":"65.872MB"}',
+                "performance": '[{"metricsKey":"cpuUsage","collectDate":"2022-12-14 12:02:05.009","metricsValue":{"cpuUsagePercentage":0.16666666666667052}},{"metricsKey":"memoryUsage","collectDate":"2022-12-14 12:02:05.009","metricsValue":{"init":65011712,"systemMaxLimit":9223372036854771712,"committed":124256256,"memUsagePercentage":9.841756766183035,"max":917504000,"used":90298512}},{"metricsKey":"memoryNoHeapUsage","collectDate":"2022-12-14 12:02:05.009","metricsValue":{"init":4121784320,"systemMaxLimit":9223372036854771712,"committed":3886866432,"memUsagePercentage":94.30057786235646,"max":4121784320,"used":3886866432}},{"metricsKey":"garbageInfo","collectDate":"2022-12-14 12:02:05.009","metricsValue":{"collectionInfoList":[{"collectionTime":18049,"collectionCount":2715,"tenured":false,"collectionName":"PS Scavenge"},{"collectionTime":948,"collectionCount":4,"tenured":true,"collectionName":"PS MarkSweep"}]}},{"metricsKey":"threadInfo","collectDate":"2022-12-14 12:02:05.225","metricsValue":{"dongTaiThreadInfoList":[{"cpuUsage":0,"cpuTime":29161574907,"name":"DongTai-IAST-ServerConfigMonitor","id":14},{"cpuUsage":0,"cpuTime":27679480672,"name":"DongTai-IAST-ConfigMonitor","id":15},{"cpuUsage":0.035035544200697016,"cpuTime":13643355644,"name":"DongTai-IAST-PerformanceMonitor","id":16},{"cpuUsage":0,"cpuTime":111735614230,"name":"DongTai-IAST-EngineMonitor","id":17},{"cpuUsage":0,"cpuTime":61038194468,"name":"DongTai-IAST-HearBeatMonitor","id":18},{"cpuUsage":0,"cpuTime":303498408,"name":"DongTai-IAST-SecondFallbackMonitor","id":19},{"cpuUsage":0,"cpuTime":859911477,"name":"DongTai-IAST-DongTaiThreadMonitor","id":20},{"cpuUsage":0,"cpuTime":26344322815,"name":"DongTai-IAST-HeartBeat","id":64}],"threadCount":51,"dongTaiThreadCount":8,"daemonThreadCount":47,"peakThreadCount":53}}]',
                 #                'isCoreRunning': 0,
-                "returnQueue": 0,
-                "cpu": "{"rate":0}"
+                "returnQueue": 0,
+                "cpu": "{"rate":0}",
             },
-            "type": 1
+            "type": 1,
         }
-        data["detail"]["agentId"] = self.agent_id
+        data["detail"]["agentId"] = self.agent_id
         res = self.agent_report(data, agentId=self.agent_id)
-        assert is_alive(self.agent_id, int(time()))
+        self.assertTrue(is_alive(self.agent_id, int(time())))
+        self.assertEqual(
+            IastHeartbeat.objects.get(agent_id=self.agent_id).cpu,
+            data["detail"]["cpu"],
+        )
+        self.assertEqual(
+            IastHeartbeat.objects.get(agent_id=self.agent_id).memory,
+            data["detail"]["memory"],
+        )
+        self.assertEqual(
+            IastHeartbeat.objects.get(agent_id=self.agent_id).disk,
+            data["detail"]["disk"],
+        )