单元测试方法-processlist

sql/engines/pgsql.py

@@ -364,7 +364,7 @@ def close(self):
 
     def processlist(self, command_type, **kwargs):
         """获取连接信息"""
-        sql ="""
+        sql = """
             select psa.pid
                                 ,concat('{',array_to_string(pg_blocking_pids(psa.pid),','),'}') block_pids
                                 ,psa.leader_pid
@@ -396,11 +396,10 @@ def processlist(self, command_type, **kwargs):
         command_type = self.escape_string(command_type)
         if not command_type:
             command_type = "Not Idle"
-        
+
         if command_type == "Not Idle":
-            sql=sql.replace("$state_not_idle$","and psa.state<>'idle'")
+            sql = sql.replace("$state_not_idle$", "and psa.state<>'idle'")
 
-        #所有的模板进行替换
-        sql=sql.replace("$state_not_idle$","")
+        # 所有的模板进行替换
+        sql = sql.replace("$state_not_idle$", "")
         return self.query("postgres", sql)
-

sql/engines/tests.py

@@ -1,6 +1,6 @@
 import json
 from datetime import timedelta, datetime
-from unittest.mock import patch, Mock, ANY
+from unittest.mock import MagicMock, patch, Mock, ANY
 
 import sqlparse
 from django.contrib.auth import get_user_model
@@ -789,6 +789,40 @@ def test_execute_workflow_exception(self, _conn, _cursor, _execute):
                 execute_result.rows[0].__dict__.keys(), row.__dict__.keys()
             )
 
+    @patch("psycopg2.connect")
+    def test_processlist_not_idle(self, mock_connect):
+        # 模拟数据库连接和游标
+        mock_cursor = MagicMock()
+        mock_connect.return_value.cursor.return_value = mock_cursor
+
+        # 假设 query 方法返回的结果
+        mock_cursor.fetchall.return_value = [
+            (123, "test_db", "user", "app_name", "active")
+        ]
+
+        # 创建 PgSQLEngine 实例
+        new_engine = PgSQLEngine(instance=self.ins)
+
+        # 调用 processlist 方法
+        result = new_engine.processlist(command_type="Not Idle")
+        self.assertEqual(result.rows, mock_cursor.fetchall.return_value)
+
+    @patch("psycopg2.connect")
+    def test_processlist_idle(self, mock_connect):
+        # 模拟数据库连接和游标
+        mock_cursor = MagicMock()
+        mock_connect.return_value.cursor.return_value = mock_cursor
+
+        # 假设 query 方法返回的结果
+        mock_cursor.fetchall.return_value = [
+            (123, "test_db", "user", "app_name", "idle")
+        ]
+        # 创建 PgSQLEngine 实例
+        new_engine = PgSQLEngine(instance=self.ins)
+        # 调用 processlist 方法
+        result = new_engine.processlist(command_type="Idle")
+        self.assertEqual(result.rows, mock_cursor.fetchall.return_value)
+
 
 class TestModel(TestCase):
     def setUp(self):