From 3001ee36e6539393f8288ad3fe898caa9803c6a9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=8E=8B=E9=A3=9E?= <feiazifeiazi@hotmail.com>
Date: Fri, 20 Sep 2024 11:27:04 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E5=8D=95=E5=85=83=E6=B5=8B?=
 =?UTF-8?q?=E8=AF=95=E6=96=B9=E6=B3=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 sql/engines/tests.py | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/sql/engines/tests.py b/sql/engines/tests.py
index c3ac4e964e..41eb142f79 100644
--- a/sql/engines/tests.py
+++ b/sql/engines/tests.py
@@ -1840,16 +1840,24 @@ def test_fill_query_columns(self):
 
     @patch("sql.engines.mongo.MongoEngine.get_connection")
     def test_current_op(self, mock_get_connection):
-        class Aggregate:
-            def __enter__(self):
-                yield {"client": "single_client"}
-                yield {"clientMetadata": {"mongos": {"client": "sharding_client"}}}
-
-            def __exit__(self, *arg, **kwargs):
-                pass
-
         mock_conn = Mock()
-        mock_conn.admin.aggregate.return_value = Aggregate()
+        # 模拟 cursor 返回的 operation 列表
+        mock_cursor = [
+            {
+                "client": "single_client",
+                "effectiveUsers": [{"user": "user_1"}],
+                "clientMetadata": {"mongos": {"client": "sharding_client"}},
+            },
+            {
+                "clientMetadata": {"mongos": {}},
+                "effectiveUsers": [{"user": "user_2"}],
+            },
+            {
+                "effectiveUsers": [],
+            },
+        ]
+
+        mock_conn.admin.aggregate.return_value = mock_cursor
         mock_get_connection.return_value = mock_conn
         command_types = ["Full", "All", "Inner", "Active"]
         for command_type in command_types: