Merge branch 'repl_set_connection'

Rakefile

@@ -38,9 +38,8 @@ end
 
 desc "Test the MongoDB Ruby driver."
 task :test do
-  puts "\nThis option has changed."
+  puts "\nTo test the driver with the C-extensions:\nrake test:c\n\n"
-  puts "\nTo test the driver with the c-extensions:\nrake test:c\n"
+  puts "To test the pure ruby driver: \nrake test:ruby\n\n"
-  puts "To test the pure ruby driver: \nrake test:ruby"
 end
 
 namespace :test do
@@ -48,22 +47,35 @@ namespace :test do
   desc "Test the driver with the C extension enabled."
   task :c do
     ENV['C_EXT'] = 'TRUE'
-    Rake::Task['test:unit'].invoke
+    if ENV['TEST']
-    Rake::Task['test:functional'].invoke
+      Rake::Task['test:functional'].invoke
-    Rake::Task['test:bson'].invoke
+    else
-    Rake::Task['test:pooled_threading'].invoke
+      Rake::Task['test:unit'].invoke
-    Rake::Task['test:drop_databases'].invoke
+      Rake::Task['test:functional'].invoke
+      Rake::Task['test:bson'].invoke
+      Rake::Task['test:pooled_threading'].invoke
+      Rake::Task['test:drop_databases'].invoke
+    end
     ENV['C_EXT'] = nil
   end
 
   desc "Test the driver using pure ruby (no C extension)"
   task :ruby do
     ENV['C_EXT'] = nil
-    Rake::Task['test:unit'].invoke
+    if ENV['TEST']
-    Rake::Task['test:functional'].invoke
+      Rake::Task['test:functional'].invoke
-    Rake::Task['test:bson'].invoke
+    else
-    Rake::Task['test:pooled_threading'].invoke
+      Rake::Task['test:unit'].invoke
-    Rake::Task['test:drop_databases'].invoke
+      Rake::Task['test:functional'].invoke
+      Rake::Task['test:bson'].invoke
+      Rake::Task['test:pooled_threading'].invoke
+      Rake::Task['test:drop_databases'].invoke
+    end
+  end
+
+  Rake::TestTask.new(:rs) do |t|
+    t.test_files = FileList['test/replica_sets/*_test.rb']
+    t.verbose    = true
   end
 
   Rake::TestTask.new(:unit) do |t|

lib/mongo.rb

@@ -42,9 +42,11 @@ require 'mongo/util/support'
 require 'mongo/util/core_ext'
 require 'mongo/util/pool'
 require 'mongo/util/server_version'
+require 'mongo/util/uri_parser'
 
 require 'mongo/collection'
 require 'mongo/connection'
+require 'mongo/repl_set_connection'
 require 'mongo/cursor'
 require 'mongo/db'
 require 'mongo/exceptions'

lib/mongo/connection.rb

@@ -35,11 +35,8 @@ module Mongo
     STANDARD_HEADER_SIZE = 16
     RESPONSE_HEADER_SIZE = 20
 
-    MONGODB_URI_MATCHER = /(([-_.\w\d]+):([-_\w\d]+)@)?([-.\w\d]+)(:([\w\d]+))?(\/([-\d\w]+))?/
-    MONGODB_URI_SPEC = "mongodb://[username:password@]host1[:port1][,host2[:port2],...[,hostN[:portN]]][/database]"
-
     attr_reader :logger, :size, :nodes, :auths, :primary, :secondaries, :arbiters,
-      :safe, :primary_pool, :read_pool, :secondary_pools
+      :safe, :primary_pool, :read_pool, :secondary_pools, :host_to_try
 
     # Counter for generating unique request ids.
     @@current_request_id = 0
@@ -92,62 +89,19 @@ module Mongo
     #
     # @core connections
     def initialize(host=nil, port=nil, options={})
-      @auths        = []
+      @host_to_try = format_pair(host, port)
-
-      if block_given?
-        @nodes = yield self
-      else
-        @nodes = format_pair(host, port)
-      end
 
       # Host and port of current master.
       @host = @port = nil
 
-      # Replica set name
-      @replica_set_name = options[:rs_name]
-
-      # Lock for request ids.
-      @id_lock = Mutex.new
-
-      # Pool size and timeout.
-      @pool_size = options[:pool_size] || 1
-      @timeout   = options[:timeout]   || 5.0
-
-      # Mutex for synchronizing pool access
-      @connection_mutex = Mutex.new
-
-      # Global safe option. This is false by default.
-      @safe = options[:safe] || false
-
-      # Create a mutex when a new key, in this case a socket,
-      # is added to the hash.
-      @safe_mutexes = Hash.new { |h, k| h[k] = Mutex.new }
-
-      # Condition variable for signal and wait
-      @queue = ConditionVariable.new
-
       # slave_ok can be true only if one node is specified
       @slave_ok = options[:slave_ok]
 
-      # Cache the various node types
+      setup(options)
-      # when connecting to a replica set.
-      @primary     = nil
-      @secondaries = []
-      @arbiters    = []
-
-      # Connection pool for primay node
-      @primary_pool    = nil
-
-      # Connection pools for each secondary node
-      @secondary_pools = []
-      @read_pool = nil
-
-      @logger   = options[:logger] || nil
-
-      should_connect = options.fetch(:connect, true)
-      connect if should_connect
     end
 
+    # DEPRECATED
+    #
     # Initialize a connection to a MongoDB replica set using an array of seed nodes.
     #
     # The seed nodes specified will be used on the initial connection to the replica set, but note
@@ -170,20 +124,13 @@ module Mongo
     #                   :read_secondary => true)
     #
     # @return [Mongo::Connection]
+    #
+    # @deprecated
     def self.multi(nodes, opts={})
-      unless nodes.length > 0 && nodes.all? {|n| n.is_a? Array}
+      warn "Connection.multi is now deprecated. Please use ReplSetConnection.new instead."
-        raise MongoArgumentError, "Connection.multi requires at least one node to be specified."
-      end
 
-      # Block returns an array, the first element being an array of nodes and the second an array
+      nodes << opts
-      # of authorizations for the database.
+      ReplSetConnection.new(*nodes)
-      new(nil, nil, opts) do |con|
-        nodes.map do |node|
-          con.instance_variable_set(:@replica_set, true)
-          con.instance_variable_set(:@read_secondary, true) if opts[:read_secondary]
-          con.pair_val_to_connection(node)
-        end
-      end
     end
 
     # Initialize a connection to MongoDB using the MongoDB URI spec:
@@ -195,8 +142,15 @@ module Mongo
     #
     # @return [Mongo::Connection]
     def self.from_uri(uri, opts={})
-      new(nil, nil, opts) do |con|
+      nodes, auths = Mongo::URIParser.parse(uri)
-        con.parse_uri(uri)
+      opts.merge!({:auths => auths})
+      if nodes.length == 1
+        Connection.new(nodes[0][0], nodes[0][1], opts)
+      elsif nodes.length > 1
+        nodes << opts
+        ReplSetConnection.new(*nodes)
+      else
+        raise MongoArgumentError, "No nodes specified. Please ensure that you've provided at least one node."
       end
     end
 
@@ -482,21 +436,12 @@ module Mongo
     # @raise [ConnectionFailure] if unable to connect to any host or port.
     def connect
       reset_connection
-      @nodes_to_try = @nodes.clone
 
-      while connecting?
+      config = check_is_master(@host_to_try)
-        node   = @nodes_to_try.shift
+      if is_primary?(config)
-        config = check_is_master(node)
+        set_primary(@host_to_try)
-
-        if is_primary?(config)
-          set_primary(node)
-        else
-          set_auxillary(node, config)
-        end
       end
 
-      pick_secondary_for_read if @read_secondary
-
       raise ConnectionFailure, "failed to connect to any given host:port" unless connected?
     end
 
@@ -515,84 +460,6 @@ module Mongo
     def close
       @primary_pool.close if @primary_pool
       @primary_pool = nil
-      @read_pool    = nil
-      @secondary_pools.each do |pool|
-        pool.close
-      end
-    end
-
-    ## Configuration helper methods
-
-    # Returns an array of host-port pairs.
-    #
-    # @private
-    def format_pair(pair_or_host, port)
-      case pair_or_host
-        when String
-          [[pair_or_host, port ? port.to_i : DEFAULT_PORT]]
-        when nil
-          [['localhost', DEFAULT_PORT]]
-      end
-    end
-
-    # Convert an argument containing a host name string and a
-    # port number integer into a [host, port] pair array.
-    #
-    # @private
-    def pair_val_to_connection(a)
-      case a
-      when nil
-        ['localhost', DEFAULT_PORT]
-      when String
-        [a, DEFAULT_PORT]
-      when Integer
-        ['localhost', a]
-      when Array
-        a
-      end
-    end
-
-    # Parse a MongoDB URI. This method is used by Connection.from_uri.
-    # Returns an array of nodes and an array of db authorizations, if applicable.
-    #
-    # @private
-    def parse_uri(string)
-      if string =~ /^mongodb:\/\//
-        string = string[10..-1]
-      else
-        raise MongoArgumentError, "MongoDB URI must match this spec: #{MONGODB_URI_SPEC}"
-      end
-
-      nodes = []
-      auths = []
-      specs = string.split(',')
-      specs.each do |spec|
-        matches  = MONGODB_URI_MATCHER.match(spec)
-        if !matches
-          raise MongoArgumentError, "MongoDB URI must match this spec: #{MONGODB_URI_SPEC}"
-        end
-
-        uname = matches[2]
-        pwd   = matches[3]
-        host  = matches[4]
-        port  = matches[6] || DEFAULT_PORT
-        if !(port.to_s =~ /^\d+$/)
-          raise MongoArgumentError, "Invalid port #{port}; port must be specified as digits."
-        end
-        port  = port.to_i
-        db    = matches[8]
-
-        if uname && pwd && db
-          add_auth(db, uname, pwd)
-        elsif uname || pwd || db
-          raise MongoArgumentError, "MongoDB URI must include all three of username, password, " +
-            "and db if any one of these is specified."
-        end
-
-        nodes << [host, port]
-      end
-
-      nodes
     end
 
     # Checkout a socket for reading (i.e., a secondary node).
@@ -629,25 +496,85 @@ module Mongo
       end
     end
 
-    private
+    protected
 
-    # Pick a node randomly from the set of possible secondaries.
+    # Generic initialization code.
-    def pick_secondary_for_read
+    # @protected
-      if (size = @secondary_pools.size) > 0
+    def setup(options)
-        @read_pool = @secondary_pools[rand(size)]
+      # Authentication objects
+      @auths = options.fetch(:auths, [])
+
+      # Lock for request ids.
+      @id_lock = Mutex.new
+
+      # Pool size and timeout.
+      @pool_size = options[:pool_size] || 1
+      @timeout   = options[:timeout]   || 5.0
+
+      # Mutex for synchronizing pool access
+      @connection_mutex = Mutex.new
+
+      # Global safe option. This is false by default.
+      @safe = options[:safe] || false
+
+      # Create a mutex when a new key, in this case a socket,
+      # is added to the hash.
+      @safe_mutexes = Hash.new { |h, k| h[k] = Mutex.new }
+
+      # Condition variable for signal and wait
+      @queue = ConditionVariable.new
+
+      # Connection pool for primay node
+      @primary      = nil
+      @primary_pool = nil
+
+      @logger   = options[:logger] || nil
+
+      should_connect = options.fetch(:connect, true)
+      connect if should_connect
+    end
+
+    ## Configuration helper methods
+
+    # Returns a host-port pair.
+    #
+    # @return [Array]
+    #
+    # @private
+    def format_pair(host, port)
+      case host
+        when String
+          [host, port ? port.to_i : DEFAULT_PORT]
+        when nil
+          ['localhost', DEFAULT_PORT]
       end
     end
 
+    # Convert an argument containing a host name string and a
+    # port number integer into a [host, port] pair array.
+    #
+    # @private
+    def pair_val_to_connection(a)
+      case a
+      when nil
+        ['localhost', DEFAULT_PORT]
+      when String
+        [a, DEFAULT_PORT]
+      when Integer
+        ['localhost', a]
+      when Array
+        a
+      end
+    end
+
+    private
+
     # If a ConnectionFailure is raised, this method will be called
     # to close the connection and reset connection values.
+    # TODO: evaluate whether this method is actually necessary
     def reset_connection
       close
       @primary = nil
-      @secondaries     = []
-      @secondary_pools = []
-      @arbiters        = []
-      @nodes_tried     = []
-      @nodes_to_try    = []
     end
 
     # Primary is defined as either a master node or a slave if
@@ -655,8 +582,9 @@ module Mongo
     #
     # If a primary node is discovered, we set the the @host and @port and
     # apply any saved authentication.
+    # TODO: simplify
     def is_primary?(config)
-      config && (config['ismaster'] == 1 || config['ismaster'] == true) || !@replica_set && @slave_ok
+      config && (config['ismaster'] == 1 || config['ismaster'] == true) || @slave_ok
     end
 
     def check_is_master(node)
@@ -666,41 +594,15 @@ module Mongo
         socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
 
         config = self['admin'].command({:ismaster => 1}, :sock => socket)
-
-        check_set_name(config, socket)
       rescue OperationFailure, SocketError, SystemCallError, IOError => ex
-        close unless connected?
+        close
       ensure
-        @nodes_tried << node
-        if config
-          update_node_list(config['hosts']) if config['hosts']
-
-          if config['msg'] && @logger
-            @logger.warn("MONGODB #{config['msg']}")
-          end
-        end
-
         socket.close if socket
       end
 
       config
     end
 
-    # Make sure that we're connected to the expected replica set.
-    def check_set_name(config, socket)
-      if @replica_set_name
-        config = self['admin'].command({:replSetGetStatus => 1},
-                   :sock => socket, :check_response => false)
-
-        if !Mongo::Support.ok?(config)
-          raise ReplicaSetConnectionError, config['errmsg']
-        elsif config['set'] != @replica_set_name
-          raise ReplicaSetConnectionError,
-            "Attempting to connect to replica set '#{config['set']}' but expected '#{@replica_set_name}'"
-        end
-      end
-    end
-
     # Set the specified node as primary, and
     # apply any saved authentication credentials.
     def set_primary(node)
@@ -710,45 +612,6 @@ module Mongo
       apply_saved_authentication
     end
 
-    # Determines what kind of node we have and caches its host
-    # and port so that users can easily connect manually.
-    def set_auxillary(node, config)
-      if config
-        if config['secondary']
-          host, port = *node
-          @secondaries << node unless @secondaries.include?(node)
-          @secondary_pools << Pool.new(self, host, port, :size => @pool_size, :timeout => @timeout)
-        elsif config['arbiterOnly']
-          @arbiters << node unless @arbiters.include?(node)
-        end
-      end
-    end
-
-    # Update the list of known nodes. Only applies to replica sets,
-    # where the response to the ismaster command will return a list
-    # of known hosts.
-    #
-    # @param hosts [Array] a list of hosts, specified as string-encoded
-    #   host-port values. Example: ["myserver-1.org:27017", "myserver-1.org:27017"]
-    #
-    # @return [Array] the updated list of nodes
-    def update_node_list(hosts)
-      new_nodes = hosts.map do |host|
-        if !host.respond_to?(:split)
-          warn "Could not parse host #{host.inspect}."
-          next
-        end
-
-        host, port = host.split(':')
-        [host, port.to_i]
-      end
-
-      # Replace the list of seed nodes with the canonical list.
-      @nodes = new_nodes.clone
-
-      @nodes_to_try = new_nodes - @nodes_tried
-    end
-
     def receive(sock, expected_response)
       begin
       receive_header(sock, expected_response)

lib/mongo/repl_set_connection.rb

@@ -0,0 +1,216 @@
+# encoding: UTF-8
+
+# --
+# Copyright (C) 2008-2010 10gen Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ++
+
+module Mongo
+
+  # Instantiates and manages connections to MongoDB.
+  class ReplSetConnection < Connection
+    attr_reader :nodes, :secondaries, :arbiters, :read_pool, :secondary_pools
+
+    def initialize(*args)
+      if args.last.is_a?(Hash)
+        opts = args.pop
+      else
+        opts = {}
+      end
+
+      unless args.length > 0
+        raise MongoArgumentError, "A ReplSetConnection requires at least one node."
+      end
+
+      # Get seed nodes
+      @nodes = args
+
+      # Replica set name
+      @replica_set = opts[:rs_name]
+
+      # Cache the various node types when connecting to a replica set.
+      @secondaries = []
+      @arbiters    = []
+
+      # Connection pools for each secondary node
+      @secondary_pools = []
+      @read_pool = nil
+
+      # Are we allowing reads from secondaries?
+      @read_secondary = opts.fetch(:read_secondary, false)
+
+      setup(opts)
+    end
+
+    # Create a new socket and attempt to connect to master.
+    # If successful, sets host and port to master and returns the socket.
+    #
+    # If connecting to a replica set, this method will replace the
+    # initially-provided seed list with any nodes known to the set.
+    #
+    # @raise [ConnectionFailure] if unable to connect to any host or port.
+    def connect
+      reset_connection
+      @nodes_to_try = @nodes.clone
+
+      while connecting?
+        node   = @nodes_to_try.shift
+        config = check_is_master(node)
+
+        if is_primary?(config)
+          set_primary(node)
+        else
+          set_auxillary(node, config)
+        end
+      end
+
+      pick_secondary_for_read if @read_secondary
+
+      if !connected?
+        if @secondary_pools.empty?
+          raise ConnectionFailure, "Failed to connect any given host:port"
+        else
+          raise ConnectionFailure, "Failed to connect to primary node."
+        end
+      end
+    end
+
+    def connecting?
+      @nodes_to_try.length > 0
+    end
+
+    # Close the connection to the database.
+    def close
+      super
+      @read_pool = nil
+      @secondary_pools.each do |pool|
+        pool.close
+      end
+    end
+
+    # If a ConnectionFailure is raised, this method will be called
+    # to close the connection and reset connection values.
+    # TODO: what's the point of this method?
+    def reset_connection
+      super
+      @secondaries     = []
+      @secondary_pools = []
+      @arbiters        = []
+      @nodes_tried  = []
+      @nodes_to_try = []
+    end
+
+    private
+
+    def check_is_master(node)
+      begin
+        host, port = *node
+        socket = TCPSocket.new(host, port)
+        socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+
+        config = self['admin'].command({:ismaster => 1}, :sock => socket)
+
+        check_set_name(config, socket)
+      rescue OperationFailure, SocketError, SystemCallError, IOError => ex
+        close unless connected?
+      ensure
+        @nodes_tried << node
+        if config
+          nodes = []
+          nodes += config['hosts'] if config['hosts']
+          nodes += config['arbiters'] if config['arbiters']
+          nodes += config['passives'] if config['passives']
+          update_node_list(nodes)
+
+          if config['msg'] && @logger
+            @logger.warn("MONGODB #{config['msg']}")
+          end
+        end
+
+        socket.close if socket
+      end
+
+      config
+    end
+
+    # Primary, when connecting to a replica can, can only be a true primary node.
+    # (And not a slave, which is possible when connecting with the standard
+    # Connection class.
+    def is_primary?(config)
+      config && (config['ismaster'] == 1 || config['ismaster'] == true)
+    end
+
+    # Pick a node randomly from the set of possible secondaries.
+    def pick_secondary_for_read
+      if (size = @secondary_pools.size) > 0
+        @read_pool = @secondary_pools[rand(size)]
+      end
+    end
+
+    # Make sure that we're connected to the expected replica set.
+    def check_set_name(config, socket)
+      if @replica_set
+        config = self['admin'].command({:replSetGetStatus => 1},
+                   :sock => socket, :check_response => false)
+
+        if !Mongo::Support.ok?(config)
+          raise ReplicaSetConnectionError, config['errmsg']
+        elsif config['set'] != @replica_set
+          raise ReplicaSetConnectionError,
+            "Attempting to connect to replica set '#{config['set']}' but expected '#{@replica_set}'"
+        end
+      end
+    end
+
+    # Determines what kind of node we have and caches its host
+    # and port so that users can easily connect manually.
+    def set_auxillary(node, config)
+      if config
+        if config['secondary']
+          host, port = *node
+          @secondaries << node unless @secondaries.include?(node)
+          @secondary_pools << Pool.new(self, host, port, :size => @pool_size, :timeout => @timeout)
+        elsif config['arbiterOnly']
+          @arbiters << node unless @arbiters.include?(node)
+        end
+      end
+    end
+
+    # Update the list of known nodes. Only applies to replica sets,
+    # where the response to the ismaster command will return a list
+    # of known hosts.
+    #
+    # @param hosts [Array] a list of hosts, specified as string-encoded
+    #   host-port values. Example: ["myserver-1.org:27017", "myserver-1.org:27017"]
+    #
+    # @return [Array] the updated list of nodes
+    def update_node_list(hosts)
+      new_nodes = hosts.map do |host|
+        if !host.respond_to?(:split)
+          warn "Could not parse host #{host.inspect}."
+          next
+        end
+
+        host, port = host.split(':')
+        [host, port.to_i]
+      end
+
+      # Replace the list of seed nodes with the canonical list.
+      @nodes = new_nodes.clone
+
+      @nodes_to_try = new_nodes - @nodes_tried
+    end
+
+  end
+end

lib/mongo/util/pool.rb

@@ -50,7 +50,12 @@ module Mongo
 
     def close
       @sockets.each do |sock|
-        sock.close
+        begin
+          sock.close
+        rescue IOError => ex
+          warn "IOError when attempting to close socket connected "
+           + "to #{@host}:#{@port}: #{ex.inspect}"
+        end
       end
       @host = @port = nil
       @sockets.clear

lib/mongo/util/uri_parser.rb

@@ -0,0 +1,71 @@
+# encoding: UTF-8
+
+# --
+# Copyright (C) 2008-2010 10gen Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ++
+
+module Mongo
+  module URIParser
+
+    DEFAULT_PORT = 27017
+    MONGODB_URI_MATCHER = /(([-_.\w\d]+):([-_\w\d]+)@)?([-.\w\d]+)(:([\w\d]+))?(\/([-\d\w]+))?/
+    MONGODB_URI_SPEC = "mongodb://[username:password@]host1[:port1][,host2[:port2],...[,hostN[:portN]]][/database]"
+
+    extend self
+
+    # Parse a MongoDB URI. This method is used by Connection.from_uri.
+    # Returns an array of nodes and an array of db authorizations, if applicable.
+    #
+    # @private
+    def parse(string)
+      if string =~ /^mongodb:\/\//
+        string = string[10..-1]
+      else
+        raise MongoArgumentError, "MongoDB URI must match this spec: #{MONGODB_URI_SPEC}"
+      end
+
+      nodes = []
+      auths = []
+      specs = string.split(',')
+      specs.each do |spec|
+        matches  = MONGODB_URI_MATCHER.match(spec)
+        if !matches
+          raise MongoArgumentError, "MongoDB URI must match this spec: #{MONGODB_URI_SPEC}"
+        end
+
+        uname = matches[2]
+        pwd   = matches[3]
+        host  = matches[4]
+        port  = matches[6] || DEFAULT_PORT
+        if !(port.to_s =~ /^\d+$/)
+          raise MongoArgumentError, "Invalid port #{port}; port must be specified as digits."
+        end
+        port  = port.to_i
+        db    = matches[8]
+
+        if uname && pwd && db
+          auths << {'db_name' => db, 'username' => uname, 'password' => pwd}
+        elsif uname || pwd || db
+          raise MongoArgumentError, "MongoDB URI must include all three of username, password, " +
+            "and db if any one of these is specified."
+        end
+
+        nodes << [host, port]
+      end
+
+      [nodes, auths]
+    end
+  end
+end

test/replica_sets/connect_test.rb

@@ -1,47 +1,78 @@
 $:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
-require 'mongo'
+require './test/replica_sets/rs_test_helper'
-require 'test/unit'
-require './test/test_helper'
 
-# NOTE: This test expects a replica set of three nodes to be running on TEST_HOST,
+# NOTE: This test expects a replica set of three nodes to be running on RS.host,
-# on ports TEST_PORT, TEST_PORT + 1, and TEST + 2.
+# on ports TEST_PORT, RS.ports[1], and TEST + 2.
 class ConnectTest < Test::Unit::TestCase
   include Mongo
 
+  def setup
+    RS.restart_killed_nodes
+  end
+
+  def teardown
+    RS.restart_killed_nodes
+  end
+
   def test_connect_bad_name
-    assert_raise_error(ReplicaSetConnectionError, "expected 'wrong-repl-set-name'") do
+    assert_raise_error(ReplicaSetConnectionError, "-wrong") do
-      Mongo::Connection.multi([[TEST_HOST, TEST_PORT], [TEST_HOST, TEST_PORT + 1], [TEST_HOST, TEST_PORT + 2]],
+      ReplSetConnection.new([RS.host, RS.ports[0]], [RS.host, RS.ports[1]],
-        :rs_name => "wrong-repl-set-name")
+        [RS.host, RS.ports[2]], :rs_name => RS.name + "-wrong")
     end
   end
 
   def test_connect
-    @conn = Mongo::Connection.multi([[TEST_HOST, TEST_PORT], [TEST_HOST, TEST_PORT + 1], [TEST_HOST, TEST_PORT + 2]],
+    @conn = ReplSetConnection.new([RS.host, RS.ports[0]], [RS.host, RS.ports[1]],
-                                    :name => "foo")
+      [RS.host, RS.ports[2]], :name => RS.name)
+    assert @conn.connected?
+
+    assert_equal RS.primary, @conn.primary
+    assert_equal RS.secondaries.sort, @conn.secondaries.sort
+    assert_equal RS.arbiters.sort, @conn.arbiters.sort
+  end
+
+  def test_connect_with_primary_node_killed
+    node = RS.kill_primary
+
+    # Becuase we're killing the primary and trying to connect right away,
+    # this is going to fail right away.
+    assert_raise_error(ConnectionFailure, "Failed to connect to primary node") do
+      @conn = ReplSetConnection.new([RS.host, RS.ports[0]], [RS.host, RS.ports[1]],
+        [RS.host, RS.ports[2]])
+    end
+
+    # This allows the secondary to come up as a primary
+    rescue_connection_failure do
+      @conn = ReplSetConnection.new([RS.host, RS.ports[0]], [RS.host, RS.ports[1]],
+        [RS.host, RS.ports[2]])
+    end
     assert @conn.connected?
   end
 
-  def test_connect_with_first_node_down
+  def test_connect_with_secondary_node_killed
-    puts "Please kill the node at #{TEST_PORT}."
+    node = RS.kill_secondary
-    gets
 
-    @conn = Mongo::Connection.multi([[TEST_HOST, TEST_PORT], [TEST_HOST, TEST_PORT + 1], [TEST_HOST, TEST_PORT + 2]])
+    @conn = ReplSetConnection.new([RS.host, RS.ports[0]], [RS.host, RS.ports[1]],
+      [RS.host, RS.ports[2]])
     assert @conn.connected?
   end
 
-  def test_connect_with_second_node_down
+  def test_connect_with_third_node_killed
-    puts "Please kill the node at #{TEST_PORT + 1}."
+    RS.kill(RS.get_node_from_port(RS.ports[2]))
-    gets
 
-    @conn = Mongo::Connection.multi([[TEST_HOST, TEST_PORT], [TEST_HOST, TEST_PORT + 1], [TEST_HOST, TEST_PORT + 2]])
+    @conn = ReplSetConnection.new([RS.host, RS.ports[0]], [RS.host, RS.ports[1]],
+      [RS.host, RS.ports[2]])
     assert @conn.connected?
   end
 
-  def test_connect_with_third_node_down
+  def test_connect_with_primary_stepped_down
-    puts "Please kill the node at #{TEST_PORT + 2}."
+    RS.step_down_primary
-    gets
 
-    @conn = Mongo::Connection.multi([[TEST_HOST, TEST_PORT], [TEST_HOST, TEST_PORT + 1], [TEST_HOST, TEST_PORT + 2]])
+    rescue_connection_failure do
+      @conn = ReplSetConnection.new([RS.host, RS.ports[0]], [RS.host, RS.ports[1]],
+        [RS.host, RS.ports[2]])
+    end
     assert @conn.connected?
   end
+
 end

test/replica_sets/count_test.rb

@@ -1,7 +1,5 @@
 $:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
-require 'mongo'
+require './test/replica_sets/rs_test_helper'
-require 'test/unit'
-require './test/test_helper'
 
 # NOTE: This test expects a replica set of three nodes to be running
 # on the local host.
@@ -9,18 +7,22 @@ class ReplicaSetCountTest < Test::Unit::TestCase
   include Mongo
 
   def setup
-    @conn = Mongo::Connection.multi([[TEST_HOST, TEST_PORT], [TEST_HOST, TEST_PORT + 1], [TEST_HOST, TEST_PORT + 2]])
+    @conn = ReplSetConnection.new([RS.host, RS.ports[0]], [RS.host, RS.ports[1]], [RS.host, RS.ports[2]])
     @db = @conn.db(MONGO_TEST_DB)
     @db.drop_collection("test-sets")
     @coll = @db.collection("test-sets")
   end
 
+  def teardown
+    RS.restart_killed_nodes
+  end
+
   def test_correct_count_after_insertion_reconnect
-    @coll.insert({:a => 20})#, :safe => {:w => 3, :wtimeout => 10000})
+    @coll.insert({:a => 20}, :safe => {:w => 2, :wtimeout => 10000})
     assert_equal 1, @coll.count
 
-    puts "Please disconnect the current master."
+    # Kill the current master node
-    gets
+    @node = RS.kill_primary
 
     rescue_connection_failure do
       @coll.insert({:a => 30}, :safe => true)

test/replica_sets/insert_test.rb

@@ -1,7 +1,5 @@
 $:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
-require 'mongo'
+require './test/replica_sets/rs_test_helper'
-require 'test/unit'
-require './test/test_helper'
 
 # NOTE: This test expects a replica set of three nodes to be running
 # on the local host.
@@ -9,16 +7,20 @@ class ReplicaSetInsertTest < Test::Unit::TestCase
   include Mongo
 
   def setup
-    @conn = Mongo::Connection.multi([[TEST_HOST, TEST_PORT], [TEST_HOST, TEST_PORT + 1], [TEST_HOST, TEST_PORT + 2]])
+    @conn = ReplSetConnection.new([TEST_HOST, RS.ports[0]], [TEST_HOST, RS.ports[1]], [TEST_HOST, RS.ports[2]])
     @db = @conn.db(MONGO_TEST_DB)
     @db.drop_collection("test-sets")
     @coll = @db.collection("test-sets")
   end
 
+  def teardown
+    RS.restart_killed_nodes
+  end
+
   def test_insert
     @coll.save({:a => 20}, :safe => true)
-    puts "Please disconnect the current master."
+
-    gets
+    RS.kill_primary
 
     rescue_connection_failure do
       @coll.save({:a => 30}, :safe => true)
@@ -29,9 +31,9 @@ class ReplicaSetInsertTest < Test::Unit::TestCase
     @coll.save({:a => 60}, :safe => true)
     @coll.save({:a => 70}, :safe => true)
 
-    puts "Please reconnect the old master to make sure that the new master " +
+    # Restart the old master and wait for sync
-         "has synced with the previous master. Note: this may have happened already."
+    RS.restart_killed_nodes
-    gets
+    sleep(1)
     results = []
 
     rescue_connection_failure do

test/replica_sets/node_type_test.rb

@@ -1,42 +0,0 @@
-$:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
-require 'mongo'
-require 'test/unit'
-require './test/test_helper'
-
-# NOTE: This test expects a replica set of three nodes, one of which is an arbiter, to be running
-# on the local host.
-class ReplicaSetNodeTypeTest < Test::Unit::TestCase
-  include Mongo
-
-  def setup
-    @conn = Mongo::Connection.multi([[TEST_HOST, TEST_PORT], [TEST_HOST, TEST_PORT + 1], [TEST_HOST, TEST_PORT + 2]])
-    @db = @conn.db(MONGO_TEST_DB)
-    @db.drop_collection("test-sets")
-    @coll = @db.collection("test-sets")
-  end
-
-  def test_correct_node_types
-    p @conn.primary
-    p @conn.secondaries
-    p @conn.arbiters
-    assert_equal 1, @conn.secondaries.length
-    assert_equal 1, @conn.arbiters.length
-
-    old_secondary = @conn.secondaries.first
-    old_primary   = @conn.primary
-
-    puts "Please disconnect the current primary and reconnect so that it becomes secondary."
-    gets
-
-    # Insert something to rescue the connection failure.
-    rescue_connection_failure do
-      @coll.insert({:a => 30}, :safe => true)
-    end
-
-    assert_equal 1, @conn.secondaries.length
-    assert_equal 1, @conn.arbiters.length
-    assert_equal old_primary, @conn.secondaries.first
-    assert_equal old_secondary, @conn.primary
-  end
-
-end

test/replica_sets/pooled_insert_test.rb

@@ -1,7 +1,5 @@
 $:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
-require 'mongo'
+require './test/replica_sets/rs_test_helper'
-require 'test/unit'
-require './test/test_helper'
 
 # NOTE: This test expects a replica set of three nodes to be running
 # on the local host.
@@ -9,18 +7,22 @@ class ReplicaSetPooledInsertTest < Test::Unit::TestCase
   include Mongo
 
   def setup
-    @conn = Mongo::Connection.multi([[TEST_HOST, TEST_PORT], [TEST_HOST, TEST_PORT + 1], [TEST_HOST, TEST_PORT + 2]],
+    @conn = ReplSetConnection.new([RS.host, RS.ports[0]], [RS.host, RS.ports[1]],
-       :pool_size => 10, :timeout => 5)
+      [RS.host, RS.ports[2]], :pool_size => 10, :timeout => 5)
     @db = @conn.db(MONGO_TEST_DB)
     @db.drop_collection("test-sets")
     @coll = @db.collection("test-sets")
   end
 
+  def teardown
+    RS.restart_killed_nodes
+  end
+
   def test_insert
     expected_results = [-1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
     @coll.save({:a => -1}, :safe => true)
-    puts "Please disconnect the current master."
+
-    gets
+    RS.kill_primary
 
     threads = []
     10.times do |i|
@@ -31,12 +33,9 @@ class ReplicaSetPooledInsertTest < Test::Unit::TestCase
       end
     end
 
-    puts "Please reconnect the old master to make sure that the new master " +
+    # Restart the old master and wait for sync
-         "has synced with the previous master. Note: this may have happened already." +
+    RS.restart_killed_nodes
-         "Note also that when connection with multiple threads, you may need to wait a few seconds" +
+    sleep(1)
-         "after restarting the old master so that all the data has had a chance to sync." +
-         "This is a case of eventual consistency."
-    gets
     results = []
 
     rescue_connection_failure do

test/replica_sets/query_secondaries.rb

@@ -1,7 +1,5 @@
 $:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
-require 'mongo'
+require './test/replica_sets/rs_test_helper'
-require 'test/unit'
-require './test/test_helper'
 
 # NOTE: This test expects a replica set of three nodes to be running
 # on the local host.
@@ -9,12 +7,16 @@ class ReplicaSetQuerySecondariesTest < Test::Unit::TestCase
   include Mongo
 
   def setup
-    @conn = Mongo::Connection.multi([[TEST_HOST, TEST_PORT]], :read_secondary => true)
+    @conn = ReplSetConnection.new([RS.host, RS.ports[0]], :read_secondary => true)
     @db = @conn.db(MONGO_TEST_DB)
     @db.drop_collection("test-sets")
     @coll = @db.collection("test-sets", :safe => {:w => 2, :wtimeout => 100})
   end
 
+  def teardown
+    RS.restart_killed_nodes
+  end
+
   def test_con
     assert @conn.primary_pool, "No primary pool!"
     assert @conn.read_pool, "No read pool!"
@@ -32,8 +34,7 @@ class ReplicaSetQuerySecondariesTest < Test::Unit::TestCase
     assert results.include?(30)
     assert results.include?(40)
 
-    puts "Please disconnect the current master."
+    RS.kill_primary
-    gets
 
     results = []
     rescue_connection_failure do

test/replica_sets/query_test.rb

@@ -1,7 +1,5 @@
 $:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
-require 'mongo'
+require './test/replica_sets/rs_test_helper'
-require 'test/unit'
-require './test/test_helper'
 
 # NOTE: This test expects a replica set of three nodes to be running
 # on the local host.
@@ -9,24 +7,27 @@ class ReplicaSetQueryTest < Test::Unit::TestCase
   include Mongo
 
   def setup
-    @conn = Mongo::Connection.multi([[TEST_HOST, TEST_PORT], [TEST_HOST, TEST_PORT + 1], [TEST_HOST, TEST_PORT + 2]])
+    @conn = ReplSetConnection.new([RS.host, RS.ports[0]])
     @db = @conn.db(MONGO_TEST_DB)
     @db.drop_collection("test-sets")
     @coll = @db.collection("test-sets")
   end
 
+  def teardown
+    RS.restart_killed_nodes
+  end
+
   def test_query
-    @coll.save({:a => 20})
+    @coll.save({:a => 20}, :safe => {:w => 3})
-    @coll.save({:a => 30})
+    @coll.save({:a => 30}, :safe => {:w => 3})
-    @coll.save({:a => 40})
+    @coll.save({:a => 40}, :safe => {:w => 3})
     results = []
     @coll.find.each {|r| results << r}
     [20, 30, 40].each do |a|
       assert results.any? {|r| r['a'] == a}, "Could not find record for a => #{a}"
     end
 
-    puts "Please disconnect the current master."
+    RS.kill_primary
-    gets
 
     results = []
     rescue_connection_failure do

test/replica_sets/replication_ack_test.rb

@@ -1,18 +1,17 @@
 $:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
-require 'mongo'
+require './test/replica_sets/rs_test_helper'
-require 'test/unit'
-require './test/test_helper'
 
 # NOTE: This test expects a replica set of three nodes to be running on local host.
 class ReplicaSetAckTest < Test::Unit::TestCase
   include Mongo
 
   def setup
-    @conn = Mongo::Connection.multi([[TEST_HOST, TEST_PORT], [TEST_HOST, TEST_PORT + 1], [TEST_HOST, TEST_PORT + 2]])
+    RS.ensure_up
 
-    master = [@conn.primary_pool.host, @conn.primary_pool.port]
+    @conn = ReplSetConnection.new([RS.host, RS.ports[0]])
 
-    @slave1 = Mongo::Connection.new(@conn.secondary_pools[0].host, @conn.secondary_pools[0].port, :slave_ok => true)
+    @slave1 = Connection.new(@conn.secondary_pools[0].host,
+      @conn.secondary_pools[0].port, :slave_ok => true)
 
     @db = @conn.db(MONGO_TEST_DB)
     @db.drop_collection("test-sets")
@@ -32,32 +31,31 @@ class ReplicaSetAckTest < Test::Unit::TestCase
   end
 
   def test_safe_mode_replication_ack
-    @col.insert({:baz => "bar"}, :safe => {:w => 2, :wtimeout => 1000})
+    @col.insert({:baz => "bar"}, :safe => {:w => 2, :wtimeout => 5000})
 
-    assert @col.insert({:foo => "0" * 10000}, :safe => {:w => 2, :wtimeout => 1000})
+    assert @col.insert({:foo => "0" * 5000}, :safe => {:w => 2, :wtimeout => 5000})
     assert_equal 2, @slave1[MONGO_TEST_DB]["test-sets"].count
 
-
+    assert @col.update({:baz => "bar"}, {:baz => "foo"}, :safe => {:w => 2, :wtimeout => 5000})
-    assert @col.update({:baz => "bar"}, {:baz => "foo"}, :safe => {:w => 2, :wtimeout => 1000})
     assert @slave1[MONGO_TEST_DB]["test-sets"].find_one({:baz => "foo"})
 
-    assert @col.remove({}, :safe => {:w => 2, :wtimeout => 1000})
+    assert @col.remove({}, :safe => {:w => 2, :wtimeout => 5000})
     assert_equal 0, @slave1[MONGO_TEST_DB]["test-sets"].count
   end
 
   def test_last_error_responses
     20.times { @col.insert({:baz => "bar"}) }
-    response = @db.get_last_error(:w => 2, :wtimeout => 10000)
+    response = @db.get_last_error(:w => 2, :wtimeout => 5000)
     assert response['ok'] == 1
     assert response['lastOp']
 
     @col.update({}, {:baz => "foo"}, :multi => true)
-    response = @db.get_last_error(:w => 2, :wtimeout => 1000)
+    response = @db.get_last_error(:w => 2, :wtimeout => 5000)
     assert response['ok'] == 1
     assert response['lastOp']
 
     @col.remove({})
-    response =  @db.get_last_error(:w => 2, :wtimeout => 1000)
+    response =  @db.get_last_error(:w => 2, :wtimeout => 5000)
     assert response['ok'] == 1
     assert response['n'] == 20
     assert response['lastOp']

test/replica_sets/rs_test_helper.rb

@@ -0,0 +1,29 @@
+$:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
+require './test/test_helper'
+require './test/tools/repl_set_manager'
+
+unless defined? RS
+  RS = ReplSetManager.new
+  RS.start_set
+end
+
+class Test::Unit::TestCase
+
+  # Generic code for rescuing connection failures and retrying operations.
+  # This could be combined with some timeout functionality.
+  def rescue_connection_failure(max_retries=60)
+    success = false
+    tries   = 0
+    while !success && tries < max_retries
+      begin
+        yield
+        success = true
+      rescue Mongo::ConnectionFailure
+        puts "Rescue attempt #{tries}\n"
+        tries += 1
+        sleep(1)
+      end
+    end
+  end
+
+end

test/test_helper.rb

@@ -69,21 +69,7 @@ class Test::Unit::TestCase
     self.class.mongo_port
   end
 
-  # Generic code for rescuing connection failures and retrying operations.
+  
-  # This could be combined with some timeout functionality.
-  def rescue_connection_failure
-    success = false
-    while !success
-      begin
-        yield
-        success = true
-      rescue Mongo::ConnectionFailure
-        puts "Rescuing"
-        sleep(1)
-      end
-    end
-  end
-
   def assert_raise_error(klass, message)
     begin
       yield

test/tools/repl_set_manager.rb

@@ -0,0 +1,241 @@
+#!/usr/bin/ruby
+
+STDOUT.sync = true
+
+unless defined? Mongo
+  require File.join(File.dirname(__FILE__), '..', '..', 'lib', 'mongo')
+end
+
+class ReplSetManager
+
+  attr_accessor :host, :start_port, :ports, :name, :mongods
+
+  def initialize(opts={})
+    @start_port = opts[:start_port] || 30000
+    @ports      = []
+    @name       = opts[:name] || 'replica-set-foo'
+    @host       = opts[:host]  || 'localhost'
+    @retries    = opts[:retries] || 60
+    @config     = {"_id" => @name, "members" => []}
+    @path       = File.join(File.expand_path(File.dirname(__FILE__)), "data")
+
+    @arbiter_count   = opts[:arbiter_count]   || 2
+    @secondary_count = opts[:secondary_count] || 1
+    @passive_count   = opts[:passive_count] || 1
+    @primary_count   = 1
+
+    @count = @primary_count + @passive_count + @arbiter_count + @secondary_count
+    if @count > 7
+      raise StandardError, "Cannot create a replica set with #{node_count} nodes. 7 is the max."
+    end
+
+    @mongods   = {}
+  end
+
+  def start_set
+    puts "** Starting a replica set with #{@count} nodes"
+
+    system("killall mongod")
+
+    n = 0
+    (@primary_count + @secondary_count).times do |n|
+      init_node(n)
+      n += 1
+    end
+
+    @passive_count.times do
+      init_node(n) do |attrs|
+        attrs['priority'] = 0
+      end
+      n += 1
+    end
+
+    @arbiter_count.times do
+      init_node(n) do |attrs|
+        attrs['arbiterOnly'] = true
+      end
+      n += 1
+    end
+
+    initiate
+    ensure_up
+  end
+
+  def init_node(n)
+    @mongods[n] ||= {}
+    port = @start_port + n
+    @ports << port
+    @mongods[n]['port'] = port
+    @mongods[n]['db_path'] = get_path("rs-#{port}")
+    @mongods[n]['log_path'] = get_path("log-#{port}")
+    system("rm -rf #{@mongods[n]['db_path']}")
+    system("mkdir -p #{@mongods[n]['db_path']}")
+
+    @mongods[n]['start'] = "mongod --replSet #{@name} --logpath '#{@mongods[n]['log_path']}' " +
+     " --dbpath #{@mongods[n]['db_path']} --port #{@mongods[n]['port']} --fork"
+
+    start(n)
+
+    member = {'_id' => n, 'host' => "#{@host}:#{@mongods[n]['port']}"}
+
+    if block_given?
+      custom_attrs = {}
+      yield custom_attrs
+      member.merge!(custom_attrs)
+      @mongods[n].merge!(custom_attrs)
+    end
+
+    @config['members'] << member
+  end
+
+  def kill(node)
+    pid = @mongods[node]['pid']
+    puts "** Killing node with pid #{pid} at port #{@mongods[node]['port']}"
+    system("kill -2 #{@mongods[node]['pid']}")
+    @mongods[node]['up'] = false
+    sleep(1)
+  end
+
+  def kill_primary
+    node = get_node_with_state(1)
+    kill(node)
+    return node
+  end
+
+  # Note that we have to rescue a connection failure
+  # when we run the StepDown command because that
+  # command will close the connection.
+  def step_down_primary
+    primary = get_node_with_state(1)
+    con = get_connection(primary)
+    begin
+      con['admin'].command({'replSetStepDown' => 90})
+    rescue Mongo::ConnectionFailure
+    end
+  end
+
+  def kill_secondary
+    node = get_node_with_state(2)
+    kill(node)
+    return node
+  end
+
+  def restart_killed_nodes
+    nodes = @mongods.keys.select do |key|
+      @mongods[key]['up'] == false
+    end
+
+    nodes.each do |node|
+      start(node)
+    end
+
+    ensure_up
+  end
+
+  def get_node_from_port(port)
+    @mongods.keys.detect { |key| @mongods[key]['port'] == port }
+  end
+
+  def start(node)
+    system(@mongods[node]['start'])
+    @mongods[node]['up'] = true
+    sleep(0.5)
+    @mongods[node]['pid'] = File.open(File.join(@mongods[node]['db_path'], 'mongod.lock')).read.strip
+  end
+  alias :restart :start
+
+  def ensure_up
+    print "** Ensuring members are up..."
+
+    attempt(Mongo::OperationFailure) do
+      con = get_connection
+      status = con['admin'].command({'replSetGetStatus' => 1})
+      print "."
+      if status['members'].all? { |m| m['health'] == 1 && [1, 2, 7].include?(m['state']) } &&
+         status['members'].any? { |m| m['state'] == 1 }
+        print "all members up!\n\n"
+        return status
+      else
+        raise Mongo::OperationFailure
+      end
+    end
+  end
+
+  def primary
+    nodes = get_all_host_pairs_with_state(1)
+    nodes.empty? ? nil : nodes[0]
+  end
+
+  def secondaries
+    get_all_host_pairs_with_state(2)
+  end
+
+  def arbiters
+    get_all_host_pairs_with_state(7)
+  end
+
+  private
+
+  def initiate
+    con = get_connection
+
+    attempt(Mongo::OperationFailure) do
+      con['admin'].command({'replSetInitiate' => @config})
+    end
+  end
+
+  def get_node_with_state(state)
+    status = ensure_up
+    node = status['members'].detect {|m| m['state'] == state}
+    if node
+      host_port = node['name'].split(':')
+      port = host_port[1] ? host_port[1].to_i : 27017
+      key = @mongods.keys.detect {|key| @mongods[key]['port'] == port}
+      return key
+    else
+      return false
+    end
+  end
+
+  def get_all_host_pairs_with_state(state)
+    status = ensure_up
+    nodes = status['members'].select {|m| m['state'] == state}
+    nodes.map do |node|
+      host_port = node['name'].split(':')
+      port = host_port[1] ? host_port[1].to_i : 27017
+      [host, port]
+    end
+  end
+
+  def get_connection(node=nil)
+    con = attempt(Mongo::ConnectionFailure) do
+      if !node
+        node = @mongods.keys.detect {|key| !@mongods[key]['arbiterOnly'] && @mongods[key]['up'] }
+      end
+      con = Mongo::Connection.new(@host, @mongods[node]['port'], :slave_ok => true)
+    end
+
+    return con
+  end
+
+  def get_path(name)
+    File.join(@path, name)
+  end
+
+  def attempt(exception)
+    raise "No block given!" unless block_given?
+    count = 0
+
+    while count < @retries do
+      begin
+        return yield
+        rescue exception
+          sleep(1)
+          count += 1
+      end
+    end
+
+    raise exception
+  end
+
+end

test/unit/connection_test.rb

@@ -41,103 +41,21 @@ class ConnectionTest < Test::Unit::TestCase
       end
     end
 
-    context "connecting to a replica set" do
-      setup do
-        TCPSocket.stubs(:new).returns(new_mock_socket('localhost', 27017))
-        @conn = Connection.multi([['localhost', 27017]], :connect => false, :read_secondary => true)
-
-        admin_db = new_mock_db
-        @hosts = ['localhost:27018', 'localhost:27019', 'localhost:27020']
-
-        admin_db.stubs(:command).returns({'ok' => 1, 'ismaster' => 1, 'hosts' => @hosts}).
-          then.returns({'ok' => 1, 'ismaster' => 0, 'hosts' => @hosts, 'secondary' => 1}).
-          then.returns({'ok' => 1, 'ismaster' => 0, 'hosts' => @hosts, 'secondary' => 1}).
-          then.returns({'ok' => 1, 'ismaster' => 0, 'arbiterOnly' => 1})
-
-        @conn.stubs(:[]).with('admin').returns(admin_db)
-        @conn.connect
-      end
-
-      should "store the hosts returned from the ismaster command" do
-        assert_equal 'localhost', @conn.primary_pool.host
-        assert_equal 27017, @conn.primary_pool.port
-
-        assert_equal 'localhost', @conn.secondary_pools[0].host
-        assert_equal 27018, @conn.secondary_pools[0].port
-
-        assert_equal 'localhost', @conn.secondary_pools[1].host
-        assert_equal 27019, @conn.secondary_pools[1].port
-
-        assert_equal 2, @conn.secondary_pools.length
-      end
-    end
-
-    context "connecting to a replica set and providing seed nodes" do
-      setup do
-        TCPSocket.stubs(:new).returns(new_mock_socket)
-        @conn = Connection.multi([['localhost', 27017], ['localhost', 27019]], :connect => false)
-
-        admin_db = new_mock_db
-        @hosts = ['localhost:27017', 'localhost:27018', 'localhost:27019']
-        admin_db.stubs(:command).returns({'ok' => 1, 'ismaster' => 1, 'hosts' => @hosts})
-        @conn.stubs(:[]).with('admin').returns(admin_db)
-        @conn.connect
-      end
-
-      should "not store any hosts redundantly" do
-      end
-    end
-
-    context "initializing a paired connection" do
-      should "require left and right nodes" do
-        assert_raise MongoArgumentError do
-          Connection.multi(['localhost', 27018], :connect => false)
-        end
-
-        assert_raise MongoArgumentError do
-          Connection.multi(['localhost', 27018], :connect => false)
-        end
-      end
-
-      should "store both nodes" do
-        @conn = Connection.multi([['localhost', 27017], ['localhost', 27018]], :connect => false)
-
-        assert_equal ['localhost', 27017], @conn.nodes[0]
-        assert_equal ['localhost', 27018], @conn.nodes[1]
-      end
-    end
-
     context "initializing with a mongodb uri" do
       should "parse a simple uri" do
         @conn = Connection.from_uri("mongodb://localhost", :connect => false)
-        assert_equal ['localhost', 27017], @conn.nodes[0]
+        assert_equal ['localhost', 27017], @conn.host_to_try
       end
 
       should "allow a complex host names" do
         host_name = "foo.bar-12345.org"
         @conn = Connection.from_uri("mongodb://#{host_name}", :connect => false)
-        assert_equal [host_name, 27017], @conn.nodes[0]
+        assert_equal [host_name, 27017], @conn.host_to_try
-      end
-
-      should "parse a uri specifying multiple nodes" do
-        @conn = Connection.from_uri("mongodb://localhost:27017,mydb.com:27018", :connect => false)
-        assert_equal ['localhost', 27017], @conn.nodes[0]
-        assert_equal ['mydb.com', 27018], @conn.nodes[1]
-      end
-
-      should "parse a uri specifying multiple nodes with auth" do
-        @conn = Connection.from_uri("mongodb://kyle:s3cr3t@localhost:27017/app,mickey:m0u5e@mydb.com:27018/dsny", :connect => false)
-        assert_equal ['localhost', 27017], @conn.nodes[0]
-        assert_equal ['mydb.com', 27018], @conn.nodes[1]
-        auth_hash = {'username' => 'kyle', 'password' => 's3cr3t', 'db_name' => 'app'}
-        assert_equal auth_hash, @conn.auths[0]
-        auth_hash = {'username' => 'mickey', 'password' => 'm0u5e', 'db_name' => 'dsny'}
-        assert_equal auth_hash, @conn.auths[1]
       end
 
       should "parse a uri with a hyphen & underscore in the username or password" do
         @conn = Connection.from_uri("mongodb://hyphen-user_name:p-s_s@localhost:27017/db", :connect => false)
-        assert_equal ['localhost', 27017], @conn.nodes[0]
+        assert_equal ['localhost', 27017], @conn.host_to_try
         auth_hash = { 'db_name' => 'db', 'username' => 'hyphen-user_name', "password" => 'p-s_s' }
         assert_equal auth_hash, @conn.auths[0]
       end

test/unit/repl_set_connection_test.rb

@@ -0,0 +1,82 @@
+require './test/test_helper'
+include Mongo
+
+class ReplSetConnectionTest < Test::Unit::TestCase
+  context "Initialization: " do
+    setup do
+      def new_mock_socket(host='localhost', port=27017)
+        socket = Object.new
+        socket.stubs(:setsockopt).with(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+        socket.stubs(:close)
+        socket
+      end
+
+      def new_mock_db
+        db = Object.new
+      end
+    end
+
+    context "connecting to a replica set" do
+      setup do
+        TCPSocket.stubs(:new).returns(new_mock_socket('localhost', 27017))
+        @conn = ReplSetConnection.new(['localhost', 27017], :connect => false, :read_secondary => true)
+
+        admin_db = new_mock_db
+        @hosts = ['localhost:27018', 'localhost:27019', 'localhost:27020']
+
+        admin_db.stubs(:command).returns({'ok' => 1, 'ismaster' => 1, 'hosts' => @hosts}).
+          then.returns({'ok' => 1, 'ismaster' => 0, 'hosts' => @hosts, 'secondary' => 1}).
+          then.returns({'ok' => 1, 'ismaster' => 0, 'hosts' => @hosts, 'secondary' => 1}).
+          then.returns({'ok' => 1, 'ismaster' => 0, 'arbiterOnly' => 1})
+
+        @conn.stubs(:[]).with('admin').returns(admin_db)
+        @conn.connect
+      end
+
+      should "store the hosts returned from the ismaster command" do
+        assert_equal 'localhost', @conn.primary_pool.host
+        assert_equal 27017, @conn.primary_pool.port
+
+        assert_equal 'localhost', @conn.secondary_pools[0].host
+        assert_equal 27018, @conn.secondary_pools[0].port
+
+        assert_equal 'localhost', @conn.secondary_pools[1].host
+        assert_equal 27019, @conn.secondary_pools[1].port
+
+        assert_equal 2, @conn.secondary_pools.length
+      end
+    end
+
+    context "connecting to a replica set and providing seed nodes" do
+      setup do
+        TCPSocket.stubs(:new).returns(new_mock_socket)
+        @conn = ReplSetConnection.new(['localhost', 27017], ['localhost', 27019], :connect => false)
+
+        admin_db = new_mock_db
+        @hosts = ['localhost:27017', 'localhost:27018', 'localhost:27019']
+        admin_db.stubs(:command).returns({'ok' => 1, 'ismaster' => 1, 'hosts' => @hosts})
+        @conn.stubs(:[]).with('admin').returns(admin_db)
+        @conn.connect
+      end
+    end
+
+    context "initializing with a mongodb uri" do
+
+      should "parse a uri specifying multiple nodes" do
+        @conn = Connection.from_uri("mongodb://localhost:27017,mydb.com:27018", :connect => false)
+        assert_equal ['localhost', 27017], @conn.nodes[0]
+        assert_equal ['mydb.com', 27018], @conn.nodes[1]
+      end
+
+      should "parse a uri specifying multiple nodes with auth" do
+        @conn = Connection.from_uri("mongodb://kyle:s3cr3t@localhost:27017/app,mickey:m0u5e@mydb.com:27018/dsny", :connect => false)
+        assert_equal ['localhost', 27017], @conn.nodes[0]
+        assert_equal ['mydb.com', 27018], @conn.nodes[1]
+        auth_hash = {'username' => 'kyle', 'password' => 's3cr3t', 'db_name' => 'app'}
+        assert_equal auth_hash, @conn.auths[0]
+        auth_hash = {'username' => 'mickey', 'password' => 'm0u5e', 'db_name' => 'dsny'}
+        assert_equal auth_hash, @conn.auths[1]
+      end
+    end
+  end
+end