RUBY-558 apply batch size on initial query

Author: estolfo

lib/mongo/cursor.rb

@@ -552,7 +552,7 @@ def construct_query_message
       message.put_int(@options)
       BSON::BSON_RUBY.serialize_cstr(message, "#{@db.name}.#{@collection.name}")
       message.put_int(@skip)
-      message.put_int(@limit)
+      @batch_size > 1 ? message.put_int(@batch_size) : message.put_int(@limit)
       spec = query_contains_special_fields? ? construct_query_spec : @selector
       message.put_binary(BSON::BSON_CODER.serialize(spec, false, false, @connection.max_bson_size).to_s)
       message.put_binary(BSON::BSON_CODER.serialize(@fields, false, false, @connection.max_bson_size).to_s) if @fields

test/functional/collection_test.rb

@@ -902,6 +902,66 @@ def test_save_symbol_find_string
     # assert_equal :mike, @@test.find_one("foo" => "mike")["foo"]
   end
 
+  def test_batch_size
+    n_docs = 6
+    batch_size = n_docs/2
+    n_docs.times do |i|
+      @@test.save(:foo => i)
+    end
+
+    doc_count = 0
+    cursor = @@test.find({}, :batch_size => batch_size)
+    cursor.next
+    assert_equal batch_size, cursor.instance_variable_get(:@returned)
+    doc_count += batch_size
+    batch_size.times { cursor.next }
+    assert_equal doc_count + batch_size, cursor.instance_variable_get(:@returned)
+    doc_count += batch_size
+    assert_equal n_docs, doc_count
+  end
+
+  def test_batch_size_with_smaller_limit
+    n_docs = 6
+    batch_size = n_docs/2
+    n_docs.times do |i|
+      @@test.insert(:foo => i)
+    end
+
+    cursor = @@test.find({}, :batch_size => batch_size, :limit => 2)
+    cursor.next
+    assert_equal 2, cursor.instance_variable_get(:@returned)
+  end
+
+  def test_batch_size_with_larger_limit
+    n_docs = 6
+    batch_size = n_docs/2
+    n_docs.times do |i|
+      @@test.insert(:foo => i)
+    end
+
+    doc_count = 0
+    cursor = @@test.find({}, :batch_size => batch_size, :limit => n_docs + 5)
+    cursor.next
+    assert_equal batch_size, cursor.instance_variable_get(:@returned)
+    doc_count += batch_size
+    batch_size.times { cursor.next }
+    assert_equal doc_count + batch_size, cursor.instance_variable_get(:@returned)
+    doc_count += batch_size
+    assert_equal n_docs, doc_count
+end
+
+  def test_batch_size_with_negative_limit
+    n_docs = 6
+    batch_size = n_docs/2
+    n_docs.times do |i|
+      @@test.insert(:foo => i)
+    end
+
+    cursor = @@test.find({}, :batch_size => batch_size, :limit => -7)
+    cursor.next
+    assert_equal n_docs, cursor.instance_variable_get(:@returned)
+  end
+
   def test_limit_and_skip
     10.times do |i|
       @@test.save(:foo => i)