From 6777a7a564ee289c44ea469ac472c9d265303bd8 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:32:50 +0800
Subject: [PATCH] =?UTF-8?q?=E8=BF=94=E5=9B=9E=E7=BB=93=E6=9E=9C=E4=BF=AE?=
 =?UTF-8?q?=E6=94=B9-=E6=A8=A1=E6=8B=9F=20MongoDB=20aggregate=20=E7=9A=84?=
 =?UTF-8?q?=E6=B8=B8=E6=A0=87=E8=A1=8C=E4=B8=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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

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