返回结果修改-模拟 MongoDB aggregate 的游标行为

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: