From c1fb8e0fb0151ab3cdf2e1bf5901a783bc194f91 Mon Sep 17 00:00:00 2001
From: guohelu <19503896967@163.com>
Date: Tue, 29 Oct 2024 15:00:00 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E8=BF=94=E5=9B=9E?=
 =?UTF-8?q?=E6=95=B0=E6=8D=AE=E6=A0=BC=E5=BC=8F=E9=94=99=E8=AF=AF=20#7550?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../utils/cmdb/test_business_host_topo.py     | 40 +++++++++++++++++--
 gcloud/utils/cmdb.py                          |  2 +-
 pipeline_plugins/cmdb_ip_picker/query.py      |  9 +++--
 3 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/gcloud/tests/utils/cmdb/test_business_host_topo.py b/gcloud/tests/utils/cmdb/test_business_host_topo.py
index d6043d00c..444797d17 100644
--- a/gcloud/tests/utils/cmdb/test_business_host_topo.py
+++ b/gcloud/tests/utils/cmdb/test_business_host_topo.py
@@ -111,7 +111,8 @@ def setUp(self):
                             },
                         ],
                     },
-                ]
+                ],
+                "count": 2,
             },
         }
         self.get_business_host_topo_expect_return = [
@@ -142,6 +143,37 @@ def setUp(self):
                 ],
             },
         ]
+        self.get_filter_business_host_topo_expect_return = (
+            [
+                {
+                    "host": {
+                        "bk_cloud_id": 0,
+                        "bk_host_id": 1,
+                        "bk_host_innerip": "127.0.0.1",
+                        "bk_mac": "",
+                        "bk_os_type": None,
+                    },
+                    "set": [{"bk_set_id": 11, "bk_set_name": "set1"}],
+                    "module": [{"bk_module_id": 56, "bk_module_name": "m1"}],
+                },
+                {
+                    "host": {
+                        "bk_cloud_id": 0,
+                        "bk_host_id": 3,
+                        "bk_host_innerip": "127.0.0.3",
+                        "bk_mac": "",
+                        "bk_os_type": None,
+                    },
+                    "set": [{"bk_set_id": 10, "bk_set_name": "空闲机池"}, {"bk_set_id": 11, "bk_set_name": "set1"}],
+                    "module": [
+                        {"bk_module_id": 54, "bk_module_name": "空闲机"},
+                        {"bk_module_id": 55, "bk_module_name": "空闲机1"},
+                        {"bk_module_id": 56, "bk_module_name": "m1"},
+                    ],
+                },
+            ],
+            2,
+        )
 
     def tearDown(self):
         self.get_client_by_user_patcher.stop()
@@ -202,7 +234,7 @@ def test__get_contains_with_ip_list(self):
             ip_str=self.ip_str,
         )
 
-        self.assertEqual(hosts_topo, self.get_business_host_topo_expect_return)
+        self.assertEqual(hosts_topo, self.get_filter_business_host_topo_expect_return)
         self.mock_client.cc.list_biz_hosts_topo.assert_called_once_with(
             {
                 "bk_biz_id": self.bk_biz_id,
@@ -228,7 +260,7 @@ def test__get_many_contains_with_ip_list(self):
             ip_str=self.ip_strs,
         )
 
-        self.assertEqual(hosts_topo, self.get_business_host_topo_expect_return)
+        self.assertEqual(hosts_topo, self.get_filter_business_host_topo_expect_return)
         self.mock_client.cc.list_biz_hosts_topo.assert_called_once_with(
             {
                 "bk_biz_id": self.bk_biz_id,
@@ -251,7 +283,7 @@ def test__get_with_page_list(self):
             self.username, self.bk_biz_id, self.supplier_account, self.host_fields, start=0, limit=10
         )
 
-        self.assertEqual(hosts_topo, self.get_business_host_topo_expect_return)
+        self.assertEqual(hosts_topo, self.get_filter_business_host_topo_expect_return)
         self.mock_client.cc.list_biz_hosts_topo.assert_called_once_with(
             {
                 "bk_biz_id": self.bk_biz_id,
diff --git a/gcloud/utils/cmdb.py b/gcloud/utils/cmdb.py
index 82953e236..ebfac5804 100644
--- a/gcloud/utils/cmdb.py
+++ b/gcloud/utils/cmdb.py
@@ -160,7 +160,7 @@ def get_filter_business_host_topo(username, bk_biz_id, supplier_account, host_fi
 
         host_info_list.append(host_info)
 
-    return host_info_list
+    return host_info_list, data["data"]["count"]
 
 
 def get_business_host(username, bk_biz_id, supplier_account, host_fields, ip_list=None, bk_cloud_id=None):
diff --git a/pipeline_plugins/cmdb_ip_picker/query.py b/pipeline_plugins/cmdb_ip_picker/query.py
index e9ef30e85..990adcadd 100644
--- a/pipeline_plugins/cmdb_ip_picker/query.py
+++ b/pipeline_plugins/cmdb_ip_picker/query.py
@@ -112,9 +112,9 @@ def cmdb_search_host(request, bk_biz_id, bk_supplier_account="", bk_supplier_id=
         message = handle_api_error(_("配置平台(CMDB)"), "cc.search_cloud_area", {}, cloud_area_result)
         result = {"result": False, "code": ERROR_CODES.API_GSE_ERROR, "message": message}
         return JsonResponse(result)
-
+    count = None
     if start and limit:
-        raw_host_info_list = cmdb.get_filter_business_host_topo(
+        raw_host_info_list, count = cmdb.get_filter_business_host_topo(
             request.user.username, bk_biz_id, bk_supplier_account, fields, int(start), int(limit), ip_str
         )
     else:
@@ -219,7 +219,10 @@ def cmdb_search_host(request, bk_biz_id, bk_supplier_account="", bk_supplier_id=
                 host_lock_status = host_lock_status_data.get(host_detail["bk_host_id"])
                 if host_lock_status is not None:
                     host_detail["bk_host_lock_status"] = host_lock_status
-    result = {"result": True, "code": NO_ERROR, "data": data}
+    if count:
+        result = {"result": True, "code": NO_ERROR, "data": {"data": data, "count": count}}
+    else:
+        result = {"result": True, "code": NO_ERROR, "data": data}
     return JsonResponse(result)