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: