模拟 MongoDB aggregate 的游标行为

hhyo · Sep 20, 2024 · e2de5ab · e2de5ab
1 parent 908f596
commit e2de5ab
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/sql/engines/tests.py b/sql/engines/tests.py
@@ -1807,13 +1807,20 @@ def test_current_op(self, mock_get_connection):
         # 模拟 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 {
+                    "client": "single_client",
+                    "effectiveUsers": [{"user": "user_1"}],
+                    "clientMetadata": {"mongos": {"client": "sharding_client"}},
+                }
+                yield {
+                    "clientMetadata": {"mongos": {}},
+                    "effectiveUsers": [{"user": "user_2"}],
+                }
                 yield {"effectiveUsers": []}
 
             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