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):