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: