PYTHON-1871 ChangeStreams must honor batchSize

Author: prashantmital

pymongo/change_stream.py

@@ -112,9 +112,11 @@ def _change_stream_options(self):
 
     def _command_options(self):
         """Return the options dict for the aggregation command."""
-        options = {'cursor': {}}
+        options = {}
         if self._max_await_time_ms is not None:
             options["maxAwaitTimeMS"] = self._max_await_time_ms
+        if self._batch_size is not None:
+            options["batchSize"] = self._batch_size
         return options
 
     def _aggregation_pipeline(self):

test/test_change_stream.py

@@ -139,6 +139,34 @@ def test_try_next_runs_one_getmore(self):
             self.assertEqual(listener.started_command_names(), ["getMore"])
             self.assertIsNone(stream.try_next())
 
+    def test_batch_size_is_honored(self):
+        listener = EventListener()
+        client = rs_or_single_client(event_listeners=[listener])
+        # Connect to the cluster.
+        client.admin.command('ping')
+        listener.results.clear()
+        # ChangeStreams only read majority committed data so use w:majority.
+        coll = self.watched_collection().with_options(
+            write_concern=WriteConcern("majority"))
+        coll.drop()
+        # Create the watched collection before starting the change stream to
+        # skip any "create" events.
+        coll.insert_one({'_id': 1})
+        self.addCleanup(coll.drop)
+        # Expected batchSize.
+        expected = {'batchSize': 23}
+        with self.change_stream_with_client(
+                client, max_await_time_ms=250, batch_size=23) as stream:
+            # Confirm that batchSize is honored for initial batch.
+            cmd = listener.results['started'][0].command
+            self.assertEqual(cmd['cursor'], expected)
+            listener.results.clear()
+            # Confirm that batchSize is honored by getMores.
+            self.assertIsNone(stream.try_next())
+            cmd = listener.results['started'][0].command
+            key = next(iter(expected))
+            self.assertEqual(expected[key], cmd[key])
+
 
 class TestClusterChangeStream(IntegrationTest, ChangeStreamTryNextMixin):