CVE-2015-4411: Uncontrolled Resource Consumption (CWE-400)
value: The Moped::BSON::ObjecId.legal? method in mongodb/bson-ruby before 3.0.4 as used in rubygem-moped allows remote attackers to cause a denial of service (worker resource consumption) via a crafted string. NOTE: This issue is due to an incomplete fix to CVE-2015-4410.
CWE description: The product does not properly control the allocation and maintenance of a limited resource, thereby enabling an actor to influence the amount of resources consumed, eventually leading to the exhaustion of available resources.
CVSS v3 base score: 7.5
Published: 2020-02-20T17:15Z
Repository: https://github.com/mongoid/moped
Language: Ruby
Files changed: .gitignore, README.md, lib/moped.rb, lib/moped/bson/object_id.rb, lib/moped/cluster.rb, lib/moped/collection.rb, lib/moped/connection.rb, lib/moped/cursor.rb, lib/moped/database.rb, lib/moped/errors.rb, lib/moped/node.rb, lib/moped/protocol/command.rb, lib/moped/query.rb, lib/moped/session.rb, lib/moped/session/context.rb, lib/moped/threaded.rb, perf/cases.rb, spec/moped/bson/object_id_spec.rb, spec/moped/cluster_spec.rb, spec/moped/collection_spec.rb, spec/moped/indexes_spec.rb, spec/moped/node_spec.rb, spec/moped/query_spec.rb, spec/moped/session_spec.rb, spec/spec_helper.rb, spec/support/mongohq.rb, spec/support/replica_set_simulator.rb, spec/support/stats.rb
Commit date: 2012-04-17 17:46:34 +0200
Commit message: Merge Replica Set Refactor

* Removes Server, and Socket; replaced with Node, and Connection.

  Replica sets are now much more robustly supported, including failover
  and discovery.

* Refactors specs.

  Internal APIs are now tested with integration specs through the public
  APIs.

* More documentation.

Diff: diff --git a/.gitignore b/.gitignore
index 2113b0b..a38b29d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,7 @@
 Gemfile.lock
 doc
 .yardoc
 .rvmrc
+.env
 perf/results
 tmp/
diff --git a/README.md b/README.md
index 38aad6b..98628b3 100644
--- a/README.md
+++ b/README.md
@@ -13,10 +13,25 @@ session[:artists].find(name: "Syd Vicious").
   update(
     :$push => { instruments: { name: "Bass" } }
   )
 ```
 
+## Features
+
+* Automated replica set node discovery and failover.
+* No C or Java extensions
+* No external dependencies
+* Simple, stable, public API.
+
+### Unsupported Features
+
+* GridFS
+* Map/Reduce
+
+These features are possible to implement, but outside the scope of Moped's
+goals. Consider them perfect opportunities to write a companion gem!
+
 # Project Breakdown
 
 Moped is composed of three parts: an implementation of the [BSON
 specification](http://bsonspec.org/), an implementation of the [Mongo Wire
 Protocol](http://www.mongodb.org/display/DOCS/Mongo+Wire+Protocol), and the
@@ -41,10 +56,35 @@ The `ObjectId` class is used for generating and interacting with Mongo's ids.
 id = Moped::BSON::ObjectId.new # => 4f8583b5e5a4e46a64000002
 id.generation_time # => 2012-04-11 13:14:29 UTC
 id == Moped::BSON::ObjectId.from_string(id.to_s) # => true
 ```
 
+<table><tbody>
+
+<tr><th>new</th>
+<td>Creates a new object id.</td></tr>
+
+<tr><th>from_string</th>
+<td>Creates a new object id from an object id string.
+<br>
+<code>Moped::BSON::ObjectId.from_string("4f8d8c66e5a4e45396000009")</code>
+</td></tr>
+
+<tr><th>from_time</th>
+<td>Creates a new object id from a time.
+<br>
+<code>Moped::BSON::ObjectId.from_time(Time.new)</code>
+</td></tr>
+
+<tr><th>legal?</th>
+<td>Validates an object id string.
+<br>
+<code>Moped::BSON::ObjectId.legal?("4f8d8c66e5a4e45396000009")</code>
+</td></tr>
+
+</tbody></table>
+
 ### Moped::BSON::Code
 
 The `Code` class is used for working with javascript on the server.
 
 ```ruby
@@ -297,10 +337,130 @@ scope.one # nil
 <td>Remove all documents matching the query's selector.<br>
 <code>query.remove_all</code></td></tr>
 
 </tbody></table>
 
+# Exceptions
+
+Here's a list of the exceptions generated by Moped.
+
+<table><tbody>
+
+<tr><th>Moped::Errors::ConnectionFailure</th>
+<td>Raised when a node cannot be reached or a connection is lost.
+<br>
+<strong>Note:</strong> this exception is only raised if Moped could not
+reconnect, so you shouldn't attempt to rescue this.</td></tr>
+
+<tr><th>Moped::Errors::OperationFailure</th>
+<td>Raised when a command fails or is invalid, such as when an insert fails in
+safe mode.</td></tr>
+
+<tr><th>Moped::Errors::QueryFailure</th>
+<td>Raised when an invalid query was sent to the database.</td></tr>
+
+<tr><th>Moped::Errors::AuthenticationFailure</th>
+<td>Raised when invalid credentials were passed to `session.login`.</td></tr>
+
+<tr><th>Moped::Errors::SocketError</th>
+<td>Not a real exception, but a module used to tag unhandled exceptions inside
+of a node's networking code. Allows you to `rescue Moped::SocketError` which
+preserving the real exception.</td></tr>
+
+</tbody></table>
+
+Other exceptions are possible while running commands, such as IO Errors around
+failed connections. Moped tries to be smart about managing its connections,
+such as checking if they're dead before executing a command; but those checks
+aren't foolproof, and Moped is conservative about handling unexpected errors on
+its connections. Namely, Moped will *not* retry a command if an unexpected
+exception is raised. Why? Because it's impossible to know whether the command
+was actually received by the remote Mongo instance, and without domain
+knowledge it cannot be safely retried.
+
+Take for example this case:
+
+```ruby
+session.with(safe: true)["users"].insert(name: "John")
+```
+
+It's entirely possible that the insert command will be sent to Mongo, but the
+connection gets closed before we read the result for `getLastError`. In this
+case, there's no way to know whether the insert was actually successful!
+
+If, however, you want to gracefully handle this in your own application, you
+could do something like:
+
+```ruby
+document = { _id: Moped::BSON::ObjectId.new, name: "John" }
+
+begin
+  session["users"].insert(document)
+rescue Moped::Errors::SocketError
+  session["users"].find(_id: document[:_id]).upsert(document)
+end
+```
+
+# Replica Sets
+
+Moped has full support for replica sets including automatic failover and node
+discovery.
+
+## Automatic Failover
+
+Moped will automatically retry lost connections and attempt to detect dead
+connections before sending an operation. Note, that it will *not* retry
+individual operations! For example, these cases will work and not raise any
+exceptions:
+
+```ruby
+session[:users].insert(name: "John")
+# kill primary node and promote secondary
+session[:users].insert(name: "John")
+session[:users].find.count # => 2.0
+
+# primary node drops our connection
+session[:users].insert(name: "John")
+```
+
+However, you'll get an operation error in a case like:
+
+```ruby
+# primary node goes down while reading the reply
+session.with(safe: true)[:users].insert(name: "John")
+```
+
+And you'll get a connection error in a case like:
+
+```ruby
+# primary node goes down, no new primary available yet
+session[:users].insert(name: "John")
+```
+
+If your session is running with eventual consistency, read operations will
+never raise connection errors as long as any secondary or primary node is
+running. The only case where you'll see a connection failure is if a node goes
+down while attempting to retrieve more results from a cursor, because cursors
+are tied to individual nodes.
+
+When two attempts to connect to a node fail, it will be marked as down. This
+removes it from the list of available nodes for `:down_interval` (default 30
+seconds). Note that the `:down_interval` only applies to normal operations;
+that is, if you ask for a primary node and none is available, all nodes will be
+retried. Likewise, if you ask for a secondary node, and no secondary or primary
+node is available, all nodes will be retreied.
+
+## Node Discovery
+
+The addresses you pass into your session are used as seeds for setting up
+replica set connections. After connection, each seed node will return a list of
+other known nodes which will be added to the set.
+
+This information is cached according to the `:refresh_interval` option (default:
+5 minutes). That means, e.g., that if you add a new node to your replica set,
+it should be represented in Moped within 5 minutes.
+
 # Thread-Safety
 
 Moped is thread-safe -- depending on your definition of thread-safe. For Moped,
 it means that there's no shared, global state between multiple sessions. What
 it doesn't mean is that a single `Session` instance can be interacted with
diff --git a/lib/moped.rb b/lib/moped.rb
index 6ca65e3..83645a5 100644
--- a/lib/moped.rb
+++ b/lib/moped.rb
@@ -4,16 +4,18 @@ require "monitor"
 require "forwardable"
 
 require "moped/bson"
 require "moped/cluster"
 require "moped/collection"
+require "moped/connection"
 require "moped/cursor"
 require "moped/database"
 require "moped/errors"
 require "moped/indexes"
 require "moped/logging"
+require "moped/node"
 require "moped/protocol"
 require "moped/query"
-require "moped/server"
 require "moped/session"
-require "moped/socket"
+require "moped/session/context"
+require "moped/threaded"
 require "moped/version"
diff --git a/lib/moped/bson/object_id.rb b/lib/moped/bson/object_id.rb
index dab72ce..6d9331a 100644
--- a/lib/moped/bson/object_id.rb
+++ b/lib/moped/bson/object_id.rb
@@ -6,117 +6,97 @@ module Moped
     class ObjectId
 
       # Formatting string for outputting an ObjectId.
       @@string_format = ("%02x" * 12).freeze
 
-      attr_reader :data
-
       class << self
         def from_string(string)
           raise Errors::InvalidObjectId.new(string) unless legal?(string)
-          data = []
+          data = ""
           12.times { |i| data << string[i*2, 2].to_i(16) }
-          new data
+          from_data data
+        end
+
+        def from_time(time)
+          from_data @@generator.generate(time.to_i)
         end
 
         def legal?(str)
-          !!str.match(/^[0-9a-f]{24}$/i)
+          !!str.match(/\A\h{24}\Z/i)
         end
-      end
 
-      def initialize(data = nil, time = nil)
-        if data
-          @data = data
-        elsif time
-          @data = @@generator.generate(time.to_i)
-        else
-          @data = @@generator.next
+        def from_data(data)
+          id = allocate
+          id.instance_variable_set :@data, data
+          id
         end
       end
 
+      def data
+        @data ||= @@generator.next
+      end
+
       def ==(other)
         BSON::ObjectId === other && data == other.data
       end
       alias eql? ==
 
       def hash
         data.hash
       end
 
       def to_s
-        @@string_format % data
+        @@string_format % data.unpack("C12")
       end
 
       # Return the UTC time at which this ObjectId was generated. This may
       # be used instread of a created_at timestamp since this information
       # is always encoded in the object id.
       def generation_time
-        Time.at(@data.pack("C4").unpack("N")[0]).utc
+        Time.at(data.unpack("N")[0]).utc
       end
 
       class << self
         def __bson_load__(io)
-          new io.read(12).unpack('C*')
+          from_data(io.read(12))
         end
-
       end
 
       def __bson_dump__(io, key)
         io << Types::OBJECT_ID
         io << key
         io << NULL_BYTE
-        io << data.pack('C12')
+        io << data
       end
 
       # @api private
       class Generator
         def initialize
           # Generate and cache 3 bytes of identifying information from the current
           # machine.
-          @machine_id = Digest::MD5.digest(Socket.gethostname).unpack("C3")
+          @machine_id = Digest::MD5.digest(Socket.gethostname).unpack("N")[0]
 
           @mutex = Mutex.new
-          @last_timestamp = nil
           @counter = 0
         end
 
-        # Return object id data based on the current time, incrementing a
-        # counter for object ids generated in the same second.
+        # Return object id data based on the current time, incrementing the
+        # object id counter.
         def next
-          now = Time.new.to_i
-
-          counter = @mutex.synchronize do
-            last_timestamp, @last_timestamp = @last_timestamp, now
-
-            if last_timestamp == now
-              @counter += 1
-            else
-              @counter = 0
-            end
+          @mutex.lock
+          begin
+            counter = @counter = (@counter + 1) % 0xFFFFFF
+          ensure
+            @mutex.unlock rescue nil
           end
 
-          generate(now, counter)
+          generate(Time.new.to_i, counter)
         end
 
-        # Generate object id data for a given time using the provided +inc+.
-        def generate(time, inc = 0)
-          pid = Process.pid % 0xFFFF
-
-          [
-            time >> 24 & 0xFF, # 4 bytes time (network order)
-            time >> 16 & 0xFF,
-            time >> 8  & 0xFF,
-            time       & 0xFF,
-            @machine_id[0],   # 3 bytes machine
-            @machine_id[1],
-            @machine_id[2],
-            pid  >> 8  & 0xFF, # 2 bytes process id
-            pid        & 0xFF,
-            inc  >> 16 & 0xFF, # 3 bytes increment
-            inc  >> 8  & 0xFF,
-            inc        & 0xFF,
-          ]
+        # Generate object id data for a given time using the provided +counter+.
+        def generate(time, counter = 0)
+          [time, @machine_id, Process.pid, counter << 8].pack("N NX lXX NX")
         end
       end
 
       @@generator = Generator.new
     end
diff --git a/lib/moped/cluster.rb b/lib/moped/cluster.rb
index bb24303..77295b3 100644
--- a/lib/moped/cluster.rb
+++ b/lib/moped/cluster.rb
@@ -1,193 +1,143 @@
 module Moped
 
-  # @api private
-  #
-  # The internal class managing connections to both a single node and replica
-  # sets.
-  #
-  # @note Though the socket class itself *is* threadsafe, the cluster presently
-  #   is not. This means that in the course of normal operations sessions can be
-  #   shared across threads, but in failure modes (when a resync is required),
-  #   things can possibly go wrong.
   class Cluster
 
-    # @return [Array] the user supplied seeds
+    # @return [Array<String>] the seeds the replica set was initialized with
     attr_reader :seeds
 
-    # @return [Boolean] whether this is a direct connection
-    attr_reader :direct
-
-    # @return [Array] all available nodes
-    attr_reader :servers
-
-    # @return [Array] seeds gathered from cluster discovery
-    attr_reader :dynamic_seeds
-
-    # @param [Array] seeds an array of host:port pairs
-    # @param [Boolean] direct (false) whether to connect directly to the hosts
-    # provided or to find additional available nodes.
-    def initialize(seeds, direct = false)
-      @seeds  = seeds
-      @direct = direct
-
-      @servers = []
-      @dynamic_seeds = []
-    end
-
-    # @return [Array] available secondary nodes
-    def secondaries
-      servers.select(&:secondary?)
-    end
-
-    # @return [Array] available primary nodes
-    def primaries
-      servers.select(&:primary?)
-    end
-
-    # @return [Array] all known addresses from user supplied seeds, dynamically
-    # discovered seeds, and active servers.
-    def known_addresses
-      [].tap do |addresses|
-        addresses.concat seeds
-        addresses.concat dynamic_seeds
-        addresses.concat servers.map { |server| server.address }
-      end.uniq
-    end
-
-    def remove(server)
-      servers.delete(server)
+    # @option options :down_interval number of seconds to wait before attempting
+    # to reconnect to a down node. (30)
+    #
+    # @option options :refresh_interval number of seconds to cache information
+    # about a node. (300)
+    def initialize(hosts, options)
+      @options = {
+        down_interval: 30,
+        refresh_interval: 300
+      }.merge(options)
+
+      @seeds = hosts
+      @nodes = hosts.map { |host| Node.new(host) }
     end
 
-    def reconnect
-      @servers = servers.map { |server| Server.new(server.address) }
-    end
-
-    def sync
-      known = known_addresses.shuffle
-      seen  = {}
-
-      sync_seed = ->(seed) do
-        server = Server.new seed
-
-        unless seen[server.resolved_address]
-          seen[server.resolved_address] = true
-
-          hosts = sync_server(server)
-
-          hosts.each do |host|
-            sync_seed[host]
-          end
+    # Refreshes information for each of the nodes provided. The node list
+    # defaults to the list of all known nodes.
+    #
+    # If a node is successfully refreshed, any newly discovered peers will also
+    # be refreshed.
+    #
+    # @return [Array<Node>] the available nodes
+    def refresh(nodes_to_refresh = @nodes)
+      refreshed_nodes = []
+      seen = {}
+
+      # Set up a recursive lambda function for refreshing a node and it's peers.
+      refresh_node = ->(node) do
+        return if seen[node]
+        seen[node] = true
+
+        # Add the node to the global list of known nodes.
+        @nodes << node unless @nodes.include?(node)
+
+        begin
+          node.refresh
+
+          # This node is good, so add it to the list of nodes to return.
+          refreshed_nodes << node unless refreshed_nodes.include?(node)
+
+          # Now refresh any newly discovered peer nodes.
+          (node.peers - @nodes).each &refresh_node
+        rescue Errors::ConnectionFailure
+          # We couldn't connect to the node, so don't do anything with it.
         end
       end
 
-      known.each do |seed|
-        sync_seed[seed]
-      end
-
-      unless servers.empty?
-        @dynamic_seeds = servers.map(&:address)
-      end
-
-      true
+      nodes_to_refresh.each &refresh_node
+      refreshed_nodes.to_a
     end
 
-    def sync_server(server)
-      [].tap do |hosts|
-        socket = server.socket
-
-        if socket.connect
-          info = socket.simple_query Protocol::Command.new(:admin, ismaster: 1)
-
-          if info["ismaster"]
-            server.primary = true
-          end
+    # Returns the list of available nodes, refreshing 1) any nodes which were
+    # down and ready to be checked again and 2) any nodes whose information is
+    # out of date.
+    #
+    # @return [Array<Node>] the list of available nodes.
+    def nodes
+      # Find the nodes that were down but are ready to be refreshed, or those
+      # with stale connection information.
+      needs_refresh, available = @nodes.partition do |node|
+        (node.down? && node.down_at < (Time.new - @options[:down_interval])) ||
+          node.needs_refresh?(Time.new - @options[:refresh_interval])
+      end
 
-          if info["secondary"]
-            server.secondary = true
-          end
+      # Refresh those nodes.
+      available.concat refresh(needs_refresh)
 
-          if info["primary"]
-            hosts.push info["primary"]
-          end
-
-          if info["hosts"]
-            hosts.concat info["hosts"]
-          end
+      # Now return all the nodes that are available.
+      available.reject &:down?
+    end
 
-          if info["passives"]
-            hosts.concat info["passives"]
+    # Yields the replica set's primary node to the provided block. This method
+    # will retry the block in case of connection errors or replica set
+    # reconfiguration.
+    #
+    # @raises ConnectionFailure when no primary node can be found
+    def with_primary(retry_on_failure = true, &block)
+      if node = nodes.find(&:primary?)
+        begin
+          node.ensure_primary do
+            return yield node.apply_auth(auth)
           end
-
-          merge(server)
-
+        rescue Errors::ConnectionFailure, Errors::ReplicaSetReconfigured
+          # Fall through to the code below if our connection was dropped or the
+          # node is no longer the primary.
         end
-      end.uniq
-    end
-
-    def merge(server)
-      previous = servers.find { |other| other == server }
-      primary = server.primary?
-      secondary = server.secondary?
+      end
 
-      if previous
-        previous.merge(server)
+      if retry_on_failure
+        # We couldn't find a primary node, so refresh the list and try again.
+        refresh
+        with_primary(false, &block)
       else
-        servers << server
+        raise Errors::ConnectionFailure, "Could not connect to a primary node for replica set #{inspect}"
       end
     end
 
-    # @param [:read, :write] mode the type of socket to return
-    # @return [Socket] a socket valid for +mode+ operations
-    def socket_for(mode)
-      sync unless primaries.any? || (secondaries.any? && mode == :read)
-
-      server = nil
-      while primaries.any? || (secondaries.any? && mode == :read)
-        if mode == :write || secondaries.empty?
-          server = primaries.sample
-        else
-          server = secondaries.sample
-        end
-
-        if server
-          socket = server.socket
-          socket.connect unless socket.connection
-
-          if socket.alive?
-            break server
-          else
-            remove server
-          end
+    # Yields a secondary node if available, otherwise the primary node. This
+    # method will retry the block in case of connection errors.
+    #
+    # @raises ConnectionError when no secondary or primary node can be found
+    def with_secondary(retry_on_failure = true, &block)
+      available_nodes = nodes.shuffle!.partition(&:secondary?).flatten
+
+      while node = available_nodes.shift
+        begin
+          return yield node.apply_auth(auth)
+        rescue Errors::ConnectionFailure
+          # That node's no good, so let's try the next one.
+          next
         end
       end
 
-      unless server
-        raise Errors::ConnectionFailure.new("Could not connect to any primary or secondary servers")
+      if retry_on_failure
+        # We couldn't find a secondary or primary node, so refresh the list and
+        # try again.
+        refresh
+        with_secondary(false, &block)
+      else
+        raise Errors::ConnectionFailure, "Could not connect to any secondary or primary nodes for replica set #{inspect}"
       end
-
-      socket = server.socket
-      socket.apply_auth auth
-      socket
     end
 
     # @return [Hash] the cached authentication credentials for this cluster.
     def auth
       @auth ||= {}
     end
 
-    # Log in to +database+ with +username+ and +password+. Does not perform the
-    # actual log in, but saves the credentials for later authentication on a
-    # socket.
-    def login(database, username, password)
-      auth[database.to_s] = [username, password]
-    end
+    private
 
-    # Log out of +database+. Does not perform the actual log out, but will log
-    # out when the socket is used next.
-    def logout(database)
-      auth.delete(database.to_s)
+    def initialize_copy(_)
+      @nodes = @nodes.map &:dup
     end
 
   end
-
 end
diff --git a/lib/moped/collection.rb b/lib/moped/collection.rb
index eac4aaf..61645a9 100644
--- a/lib/moped/collection.rb
+++ b/lib/moped/collection.rb
@@ -54,14 +54,12 @@ module Moped
     #   @example
     #     db[:people].insert([{name: "John"}, {name: "Joe"}])
     #   @param [Array<Hash>] documents the documents to insert
     def insert(documents)
       documents = [documents] unless documents.is_a? Array
-      insert = Protocol::Insert.new(database.name, name, documents)
 
       database.session.with(consistency: :strong) do |session|
-        session.execute insert
+        session.context.insert(database.name, name, documents)
       end
-
     end
   end
 end
diff --git a/lib/moped/connection.rb b/lib/moped/connection.rb
new file mode 100644
index 0000000..d94d254
--- /dev/null
+++ b/lib/moped/connection.rb
@@ -0,0 +1,95 @@
+require "timeout"
+
+module Moped
+  class Connection
+
+    class TCPSocket < ::TCPSocket
+      def self.connect(host, port, timeout)
+        Timeout::timeout(timeout) do
+          new(host, port).tap do |sock|
+            sock.set_encoding 'binary'
+          end
+        end
+      end
+
+      def alive?
+        if Kernel::select([self], nil, nil, 0)
+          !eof? rescue false
+        else
+          true
+        end
+      end
+
+      def write(*args)
+        raise Errors::ConnectionFailure, "Socket connection was closed by remote host" unless alive?
+        super
+      end
+    end
+
+    def initialize
+      @sock = nil
+      @request_id = 0
+    end
+
+    def connect(host, port, timeout)
+      @sock = TCPSocket.connect host, port, timeout
+    end
+
+    def alive?
+      connected? ? @sock.alive? : false
+    end
+
+    def connected?
+      !!@sock
+    end
+
+    def disconnect
+      @sock.close
+    rescue
+    ensure
+      @sock = nil
+    end
+
+    def write(operations)
+      buf = ""
+
+      operations.each do |operation|
+        operation.request_id = (@request_id += 1)
+        operation.serialize(buf)
+      end
+
+      @sock.write buf
+    end
+
+    def receive_replies(operations)
+      operations.map do |operation|
+        read if operation.is_a?(Protocol::Query) || operation.is_a?(Protocol::GetMore)
+      end
+    end
+
+    def read
+      reply = Protocol::Reply.allocate
+
+      reply.length,
+        reply.request_id,
+        reply.response_to,
+        reply.op_code,
+        reply.flags,
+        reply.cursor_id,
+        reply.offset,
+        reply.count = @sock.read(36).unpack('l5<q<l2<')
+
+      if reply.count == 0
+        reply.documents = []
+      else
+        buffer = StringIO.new(@sock.read(reply.length - 36))
+
+        reply.documents = reply.count.times.map do
+          BSON::Document.deserialize(buffer)
+        end
+      end
+
+      reply
+    end
+  end
+end
diff --git a/lib/moped/cursor.rb b/lib/moped/cursor.rb
index 907ed50..359403b 100644
--- a/lib/moped/cursor.rb
+++ b/lib/moped/cursor.rb
@@ -8,53 +8,72 @@ module Moped
     attr_reader :get_more_op
     attr_reader :kill_cursor_op
 
     def initialize(session, query_operation)
       @session = session
-      @query_op = query_operation.dup
 
-      @get_more_op = Protocol::GetMore.new(
-        @query_op.database,
-        @query_op.collection,
-        0,
-        @query_op.limit
-      )
+      @database    = query_operation.database
+      @collection  = query_operation.collection
+      @selector    = query_operation.selector
 
-      @kill_cursor_op = Protocol::KillCursors.new([0])
+      @cursor_id = 0
+      @limit = query_operation.limit
+      @limited = @limit > 0
+
+      @options = {
+        request_id: query_operation.request_id,
+        flags: query_operation.flags,
+        limit: query_operation.limit,
+        skip: query_operation.skip,
+        fields: query_operation.fields
+      }
     end
 
     def each
-      documents = query @query_op
+      documents = load
       documents.each { |doc| yield doc }
 
       while more?
-        return kill if limited? && @get_more_op.limit <= 0
+        return kill if limited? && @limit <= 0
 
-        documents = query @get_more_op
+        documents = get_more
         documents.each { |doc| yield doc }
       end
     end
 
-    def query(operation)
-      reply = session.query operation
+    def load
+      consistency = session.consistency
+      @options[:flags] |= [:slave_ok] if consistency == :eventual
+
+      reply, @node = session.context.with_node do |node|
+        [node.query(@database, @collection, @selector, @options), node]
+      end
 
-      @get_more_op.limit -= reply.count if limited?
-      @get_more_op.cursor_id = reply.cursor_id
-      @kill_cursor_op.cursor_ids = [reply.cursor_id]
+      @limit -= reply.count if limited?
+      @cursor_id = reply.cursor_id
 
       reply.documents
     end
 
     def limited?
-      @query_op.limit > 0
+      @limited
     end
 
     def more?
-      @get_more_op.cursor_id != 0
+      @cursor_id != 0
+    end
+
+    def get_more
+      reply = @node.get_more @database, @collection, @cursor_id, @limit
+
+      @limit -= reply.count if limited?
+      @cursor_id = reply.cursor_id
+
+      reply.documents
     end
 
     def kill
-      session.execute kill_cursor_op
+      @node.kill_cursors [@cursor_id]
     end
   end
 
 end
diff --git a/lib/moped/database.rb b/lib/moped/database.rb
index 6518abb..a69b0ed 100644
--- a/lib/moped/database.rb
+++ b/lib/moped/database.rb
@@ -27,24 +27,26 @@ module Moped
       @name = name
     end
 
     # Drop the database.
     def drop
-      command dropDatabase: 1
+      session.with(consistency: :strong) do |session|
+        session.context.command name, dropDatabase: 1
+      end
     end
 
     # Log in with +username+ and +password+ on the current database.
     #
     # @param [String] username the username
     # @param [String] password the password
     def login(username, password)
-      session.cluster.login(name, username, password)
+      session.context.login(name, username, password)
     end
 
     # Log out from the current database.
     def logout
-      session.cluster.logout(name)
+      session.context.logout(name)
     end
 
     # Run +command+ on the database.
     #
     # @example
@@ -52,21 +54,11 @@ module Moped
     #   # => { "master" => true, hosts: [] }
     #
     # @param [Hash] command the command to run
     # @return [Hash] the result of the command
     def command(command)
-      operation = Protocol::Command.new(name, command)
-
-      result = session.with(consistency: :strong) do |session|
-        session.simple_query(operation)
-      end
-
-      raise Errors::OperationFailure.new(
-        operation, result
-      ) unless result["ok"] == 1.0
-
-      result
+      session.context.command name, command
     end
 
     # @param [Symbol, String] collection the collection name
     # @return [Moped::Collection] an instance of +collection+
     def [](collection)
diff --git a/lib/moped/errors.rb b/lib/moped/errors.rb
index 6f94984..da88514 100644
--- a/lib/moped/errors.rb
+++ b/lib/moped/errors.rb
@@ -1,32 +1,21 @@
 module Moped
-
-  # The namespace for all errors generated by Moped.
   module Errors
 
     # Mongo's exceptions are sparsely documented, but this is the most accurate
     # source of information on error codes.
     ERROR_REFERENCE = "https://github.com/mongodb/mongo/blob/master/docs/errors.md"
 
     # Generic error class for exceptions related to connection failures.
     class ConnectionFailure < StandardError; end
 
-    # Raised when providing an invalid string from an object id.
-    class InvalidObjectId < StandardError
-      def initialize(string)
-        super("'#{string}' is not a valid object id.")
-      end
-    end
+    # Tag applied to unhandled exceptions on a node.
+    module SocketError end
 
     # Generic error class for exceptions generated on the remote MongoDB
     # server.
-    class MongoError < StandardError; end
-
-    # Exception class for exceptions generated as a direct result of an
-    # operation, such as a failed insert or an invalid command.
-    class OperationFailure < MongoError
-
+    class MongoError < StandardError
       # @return the command that generated the error
       attr_reader :command
 
       # @return [Hash] the details about the error
       attr_reader :details
@@ -58,11 +47,30 @@ module Moped
           "failed with error #{err.inspect}"
         end
       end
     end
 
-    # A special kind of OperationFailure, raised when Mongo sets the
-    # :query_failure flag on a query response.
-    class QueryFailure < OperationFailure; end
+    # Exception class for exceptions generated as a direct result of an
+    # operation, such as a failed insert or an invalid command.
+    class OperationFailure < MongoError; end
+
+    # Exception raised on invalid queries.
+    class QueryFailure < MongoError; end
+
+    # Exception raised when authentication fails.
+    class AuthenticationFailure < MongoError; end
+
+    # Raised when providing an invalid string from an object id.
+    class InvalidObjectId < StandardError
+      def initialize(string)
+        super("'#{string}' is not a valid object id.")
+      end
+    end
+
+    # @api private
+    #
+    # Internal exception raised by Node#ensure_primary and captured by
+    # Cluster#with_primary.
+    class ReplicaSetReconfigured < StandardError; end
 
   end
 end
diff --git a/lib/moped/node.rb b/lib/moped/node.rb
new file mode 100644
index 0000000..e9a8d57
--- /dev/null
+++ b/lib/moped/node.rb
@@ -0,0 +1,334 @@
+module Moped
+  class Node
+
+    attr_reader :address
+    attr_reader :resolved_address
+    attr_reader :ip_address
+    attr_reader :port
+
+    attr_reader :peers
+    attr_reader :timeout
+
+    def initialize(address)
+      @address = address
+
+      host, port = address.split(":")
+      @ip_address = ::Socket.getaddrinfo(host, nil, ::Socket::AF_INET, ::Socket::SOCK_STREAM).first[3]
+      @port = port.to_i
+      @resolved_address = "#{@ip_address}:#{@port}"
+
+      @timeout = 5
+    end
+
+    def command(database, cmd, options = {})
+      operation = Protocol::Command.new(database, cmd, options)
+
+      process(operation) do |reply|
+        result = reply.documents[0]
+
+        raise Errors::OperationFailure.new(
+          operation, result
+        ) if result["ok"] != 1 || result["err"] || result["errmsg"]
+
+        result
+      end
+    end
+
+    def kill_cursors(cursor_ids)
+      process Protocol::KillCursors.new(cursor_ids)
+    end
+
+    def get_more(database, collection, cursor_id, limit)
+      process Protocol::GetMore.new(database, collection, cursor_id, limit)
+    end
+
+    def remove(database, collection, selector, options = {})
+      process Protocol::Delete.new(database, collection, selector, options)
+    end
+
+    def update(database, collection, selector, change, options = {})
+      process Protocol::Update.new(database, collection, selector, change, options)
+    end
+
+    def insert(database, collection, documents)
+      process Protocol::Insert.new(database, collection, documents)
+    end
+
+    def query(database, collection, selector, options = {})
+      operation = Protocol::Query.new(database, collection, selector, options)
+
+      process operation do |reply|
+        if reply.flags.include? :query_failure
+          raise Errors::QueryFailure.new(operation, reply.documents.first)
+        end
+
+        reply
+      end
+    end
+
+    # @return [true/false] whether the node needs to be refreshed.
+    def needs_refresh?(time)
+      !@refreshed_at || @refreshed_at < time
+    end
+
+    def primary?
+      @primary
+    end
+
+    def secondary?
+      @secondary
+    end
+
+    # Refresh information about the node, such as it's status in the replica
+    # set and it's known peers.
+    #
+    # Returns nothing.
+    # Raises Errors::ConnectionFailure if the node cannot be reached
+    # Raises Errors::ReplicaSetReconfigured if the node is no longer a primary node and
+    #   refresh was called within an +#ensure_primary+ block.
+    def refresh
+      info = command "admin", ismaster: 1
+
+      @refreshed_at = Time.now
+      primary = true if info["ismaster"]
+      secondary = true if info["secondary"]
+
+      peers = []
+      peers.push info["primary"] if info["primary"]
+      peers.concat info["hosts"] if info["hosts"]
+      peers.concat info["passives"] if info["passives"]
+      peers.concat info["arbiters"] if info["arbiters"]
+
+      @peers = peers.map { |peer| Node.new(peer) }
+      @primary, @secondary = primary, secondary
+
+      if !primary && Threaded.executing?(:ensure_primary)
+        raise Errors::ReplicaSetReconfigured, "#{inspect} is no longer the primary node."
+      end
+    end
+
+    attr_reader :down_at
+
+    def down?
+      @down_at
+    end
+
+    # Set a flag on the node for the duration of provided block so that an
+    # exception is raised if the node is no longer the primary node.
+    #
+    # Returns nothing.
+    def ensure_primary
+      Threaded.begin :ensure_primary
+      yield
+    ensure
+      Threaded.end :ensure_primary
+    end
+
+    # Yields the block if a connection can be established, retrying when a
+    # connection error is raised.
+    #
+    # @raises ConnectionFailure when a connection cannot be established.
+    def ensure_connected
+      # Don't run the reconnection login if we're already inside an
+      # +ensure_connected+ block.
+      return yield if Threaded.executing? :connection
+      Threaded.begin :connection
+
+      retry_on_failure = true
+
+      begin
+        connect unless connected?
+        yield
+      rescue Errors::ReplicaSetReconfigured
+        # Someone else wrapped this in an #ensure_primary block, so let the
+        # reconfiguration exception bubble up.
+        raise
+      rescue Errors::OperationFailure, Errors::AuthenticationFailure, Errors::QueryFailure
+        # These exceptions are "expected" in the normal course of events, and
+        # don't necessitate disconnecting.
+        raise
+      rescue Errors::ConnectionFailure
+        disconnect
+
+        if retry_on_failure
+          # Maybe there was a hiccup -- try reconnecting one more time
+          retry_on_failure = false
+          retry
+        else
+          # Nope, we failed to connect twice. Flag the node as down and re-raise
+          # the exception.
+          down!
+          raise
+        end
+      rescue
+        # Looks like we got an unexpected error, so we'll clean up the connection
+        # and re-raise the exception.
+        disconnect
+        raise $!.extend(Errors::SocketError)
+      end
+    ensure
+      Threaded.end :connection
+    end
+
+    def pipeline
+      Threaded.begin :pipeline
+
+      begin
+        yield
+      ensure
+        Threaded.end :pipeline
+      end
+
+      flush unless Threaded.executing? :pipeline
+    end
+
+    def apply_auth(credentials)
+      unless auth == credentials
+        logouts = auth.keys - credentials.keys
+
+        logouts.each do |database|
+          logout database
+        end
+
+        credentials.each do |database, (username, password)|
+          login(database, username, password) unless auth[database] == [username, password]
+        end
+      end
+
+      self
+    end
+
+    def ==(other)
+      resolved_address == other.resolved_address
+    end
+    alias eql? ==
+
+    def hash
+      [ip_address, port].hash
+    end
+
+    private
+
+    def auth
+      @auth ||= {}
+    end
+
+    def login(database, username, password)
+      getnonce = Protocol::Command.new(database, getnonce: 1)
+      connection.write [getnonce]
+      result = connection.read.documents.first
+      raise Errors::OperationFailure.new(getnonce, result) unless result["ok"] == 1
+
+      authenticate = Protocol::Commands::Authenticate.new(database, username, password, result["nonce"])
+      connection.write [authenticate]
+      result = connection.read.documents.first
+      raise Errors::AuthenticationFailure.new(authenticate, result) unless result["ok"] == 1
+
+      auth[database] = [username, password]
+    end
+
+    def logout(database)
+      command = Protocol::Command.new(database, logout: 1)
+      connection.write [command]
+      result = connection.read.documents.first
+      raise Errors::OperationFailure.new(command, result) unless result["ok"] == 1
+
+      auth.delete(database)
+    end
+
+    def initialize_copy(_)
+      @connection = nil
+    end
+
+    def connection
+      @connection ||= Connection.new
+    end
+
+    def disconnect
+      auth.clear
+      connection.disconnect
+    end
+
+    def connected?
+      connection.connected?
+    end
+
+    # Mark the node as down.
+    #
+    # Returns nothing.
+    def down!
+      @down_at = Time.new
+
+      disconnect
+    end
+
+    # Connect to the node.
+    #
+    # Returns nothing.
+    # Raises Moped::ConnectionError if the connection times out.
+    # Raises Moped::ConnectionError if the server is unavailable.
+    def connect
+      connection.connect ip_address, port, timeout
+      @down_at = nil
+
+      refresh
+    rescue Timeout::Error
+      raise Errors::ConnectionFailure, "Timed out connection to Mongo on #{address}"
+    rescue Errno::ECONNREFUSED
+      raise Errors::ConnectionFailure, "Could not connect to Mongo on #{address}"
+    end
+
+    def process(operation, &callback)
+      if Threaded.executing? :pipeline
+        queue.push [operation, callback]
+      else
+        flush([[operation, callback]])
+      end
+    end
+
+    def queue
+      Threaded.stack(:pipelined_operations)
+    end
+
+    def flush(ops = queue)
+      operations, callbacks = ops.transpose
+
+      logging(operations) do
+        ensure_connected do
+          connection.write operations
+          replies = connection.receive_replies(operations)
+
+          replies.zip(callbacks).map do |reply, callback|
+            callback ? callback[reply] : reply
+          end.last
+        end
+      end
+    ensure
+      ops.clear
+    end
+
+    def logging(operations)
+      instrument_start = (logger = Moped.logger) && logger.debug? && Time.new
+      yield
+    ensure
+      log_operations(logger, operations, Time.new - instrument_start) if instrument_start && !$!
+    end
+
+    def log_operations(logger, ops, duration)
+      prefix  = "  MOPED: #{address} "
+      indent  = " "*prefix.length
+      runtime = (" (%.1fms)" % duration)
+
+      if ops.length == 1
+        logger.debug prefix + ops.first.log_inspect + runtime
+      else
+        first, *middle, last = ops
+
+        logger.debug prefix + first.log_inspect
+        middle.each { |m| logger.debug indent + m.log_inspect }
+        logger.debug indent + last.log_inspect + runtime
+      end
+    end
+
+  end
+end
diff --git a/lib/moped/protocol/command.rb b/lib/moped/protocol/command.rb
index aeaf0c4..4723b60 100644
--- a/lib/moped/protocol/command.rb
+++ b/lib/moped/protocol/command.rb
@@ -9,12 +9,13 @@ module Moped
     #   socket.write command.serialize
     class Command < Query
 
       # @param [String, Symbol] database the database to run this command on
       # @param [Hash] command the command to run
-      def initialize(database, command)
-        super database, :$cmd, command, limit: -1
+      # @param [Hash] additional query options
+      def initialize(database, command, options = {})
+        super database, :$cmd, command, options.merge(limit: -1)
       end
 
       def log_inspect
         type = "COMMAND"
 
diff --git a/lib/moped/query.rb b/lib/moped/query.rb
index b2e676a..e347e94 100644
--- a/lib/moped/query.rb
+++ b/lib/moped/query.rb
@@ -97,19 +97,19 @@ module Moped
       self
     end
 
     # @return [Hash] the first document that matches the selector.
     def first
-      session.simple_query(operation)
+      limit(-1).each.first
     end
     alias one first
 
     # Iterate through documents matching the query's selector.
     #
     # @yieldparam [Hash] document each matching document
     def each
-      cursor = Cursor.new(session.with(retain_socket: true), operation)
+      cursor = Cursor.new(session, operation)
       cursor.to_enum.tap do |enum|
         enum.each do |document|
           yield document
         end if block_given?
       end
@@ -127,10 +127,11 @@ module Moped
       result = collection.database.command(
         distinct: collection.name,
         key: key.to_s,
         query: selector
       )
+
       result["values"]
     end
 
     # @return [Numeric] the number of documents that match the selector.
     def count
@@ -149,20 +150,16 @@ module Moped
     #
     # @param [Hash] change the changes to make to the document
     # @param [Array] flags an array of operation flags. Valid values are:
     #   +:multi+ and +:upsert+
     def update(change, flags = nil)
-      update = Protocol::Update.new(
-        operation.database,
-        operation.collection,
-        operation.selector,
-        change,
-        flags: flags
-      )
-
       session.with(consistency: :strong) do |session|
-        session.execute update
+        session.context.update operation.database,
+          operation.collection,
+          operation.selector,
+          change,
+          flags: flags
       end
     end
 
     # Update multiple documents matching the query's selector.
     #
@@ -191,35 +188,27 @@ module Moped
     # Remove a single document matching the query's selector.
     #
     # @example
     #   db[:people].find(name: "John").remove
     def remove
-      delete = Protocol::Delete.new(
-        operation.database,
-        operation.collection,
-        operation.selector,
-        flags: [:remove_first]
-      )
-
       session.with(consistency: :strong) do |session|
-        session.execute delete
+        session.context.remove operation.database,
+          operation.collection,
+          operation.selector,
+          flags: [:remove_first]
       end
     end
 
     # Remove multiple documents matching the query's selector.
     #
     # @example
     #   db[:people].find(name: "John").remove_all
     def remove_all
-      delete = Protocol::Delete.new(
-        operation.database,
-        operation.collection,
-        operation.selector
-      )
-
       session.with(consistency: :strong) do |session|
-        session.execute delete
+        session.context.remove operation.database,
+          operation.collection,
+          operation.selector
       end
     end
 
     private
 
diff --git a/lib/moped/server.rb b/lib/moped/server.rb
deleted file mode 100644
index 3bdfcab..0000000
--- a/lib/moped/server.rb
+++ /dev/null
@@ -1,73 +0,0 @@
-module Moped
-
-  # @api private
-  #
-  # The internal class for storing information about a server.
-  class Server
-
-    # @return [String] the original host:port address provided
-    attr_reader :address
-
-    # @return [String] the resolved host:port address
-    attr_reader :resolved_address
-
-    # @return [String] the resolved ip address
-    attr_reader :ip_address
-
-    # @return [Integer] the resolved port
-    attr_reader :port
-
-    attr_writer :primary
-    attr_writer :secondary
-
-    def initialize(address)
-      @address = address
-
-      host, port = address.split(":")
-      port = port ? port.to_i : 27017
-
-      ip_address = ::Socket.getaddrinfo(host, nil, ::Socket::AF_INET, ::Socket::SOCK_STREAM).first[3]
-
-      @primary = @secondary = false
-      @ip_address = ip_address
-      @port = port
-      @resolved_address = "#{ip_address}:#{port}"
-    end
-
-    def primary?
-      !!@primary
-    end
-
-    def secondary?
-      !!@secondary
-    end
-
-    def merge(other)
-      @primary = other.primary?
-      @secondary = other.secondary?
-
-      other.close
-    end
-
-    def close
-      if @socket
-        @socket.close
-        @socket = nil
-      end
-    end
-
-    def socket
-      @socket ||= Socket.new(ip_address, port)
-    end
-
-    def ==(other)
-      self.class === other && hash == other.hash
-    end
-    alias eql? ==
-
-    def hash
-      [ip_address, port].hash
-    end
-
-  end
-end
diff --git a/lib/moped/session.rb b/lib/moped/session.rb
index 2b505bb..a4e9e46 100644
--- a/lib/moped/session.rb
+++ b/lib/moped/session.rb
@@ -35,33 +35,43 @@ module Moped
 
     # @private
     # @return [Cluster] this session's cluster
     attr_reader :cluster
 
+    # @private
+    # @return [Context] this session's context
+    attr_reader :context
+
     # @param [Array] seeds an of host:port pairs
     # @param [Hash] options
     # @option options [Boolean] :safe (false) ensure writes are persisted
     # @option options [Hash] :safe ensure writes are persisted with the
     #   specified safety level e.g., "fsync: true", or "w: 2, wtimeout: 5"
     # @option options [Symbol, String] :database the database to use
     # @option options [:strong, :eventual] :consistency (:eventual)
     def initialize(seeds, options = {})
-      @cluster = Cluster.new(seeds)
+      @cluster = Cluster.new(seeds, {})
+      @context = Context.new(self)
       @options = options
       @options[:consistency] ||= :eventual
     end
 
     # @return [Boolean] whether the current session requires safe operations.
     def safe?
       !!safety
     end
 
+    # @return [:strong, :eventual] the session's consistency
+    def consistency
+      options[:consistency]
+    end
+
     # Switch the session's current database.
     #
     # @example
     #   session.use :moped
-    #   session[:people].     john, mary = session[:people].find.one # => { :name => "John" }
+    #   session[:people].find.one # => { :name => "John" }
     #
     # @param [String] database the database to use
     def use(database)
       options[:database] = database
       set_current_database database
@@ -109,11 +119,11 @@ module Moped
     # @see #with
     # @yieldparam [Moped::Session] session the new session
     # @return [Moped::Session] the new session
     def new(options = {})
       session = with(options)
-      session.cluster.reconnect
+      session.instance_variable_set(:@cluster, cluster.dup)
 
       if block_given?
         yield session
       else
         session
@@ -153,68 +163,10 @@ module Moped
     #
     # @param (see Moped::Database#logout)
     # @raise (see Moped::Database#login)
     delegate :logout => :current_database
 
-    # @private
-    def current_database
-      return @current_database if defined? @current_database
-
-      if database = options[:database]
-        set_current_database(database)
-      else
-        raise "No database set for session. Call #use or #with before accessing the database"
-      end
-    end
-
-    # @private
-    def simple_query(query)
-      query.limit = -1
-
-      query(query).documents.first
-    end
-
-    # @private
-    def query(query)
-      if options[:consistency] == :eventual
-        query.flags |= [:slave_ok] if query.respond_to? :flags
-        mode = :read
-      else
-        mode = :write
-      end
-
-      reply = socket_for(mode).execute(query)
-
-      reply.tap do |reply|
-        if reply.flags.include?(:query_failure)
-          raise Errors::QueryFailure.new(query, reply.documents.first)
-        end
-      end
-    end
-
-    # @private
-    def execute(op)
-      mode = options[:consistency] == :eventual ? :read : :write
-      socket = socket_for(mode)
-
-      if safe?
-        last_error = Protocol::Command.new(
-          "admin", { getlasterror: 1 }.merge(safety)
-        )
-
-        socket.execute(op, last_error).documents.first.tap do |result|
-          raise Errors::OperationFailure.new(
-            op, result
-          ) if result["err"] || result["errmsg"]
-        end
-      else
-        socket.execute(op)
-      end
-    end
-
-    private
-
     # @return [Boolean, Hash] the safety level for this session
     def safety
       safe = options[:safe]
 
       case safe
@@ -225,29 +177,31 @@ module Moped
       else
         safe
       end
     end
 
-    def socket_for(mode)
-      if options[:retain_socket]
-        @socket ||= cluster.socket_for(mode)
+    private
+
+    def current_database
+      return @current_database if defined? @current_database
+
+      if database = options[:database]
+        set_current_database(database)
       else
-        cluster.socket_for(mode)
+        raise "No database set for session. Call #use or #with before accessing the database"
       end
     end
 
     def set_current_database(database)
       @current_database = Database.new(self, database)
     end
 
-    def dup
-      session = super
-      session.instance_variable_set :@options, options.dup
+    def initialize_copy(_)
+      @context = Context.new(self)
+      @options = @options.dup
 
       if defined? @current_database
-        session.send(:remove_instance_variable, :@current_database)
+        remove_instance_variable :@current_database
       end
-
-      session
     end
   end
 end
diff --git a/lib/moped/session/context.rb b/lib/moped/session/context.rb
new file mode 100644
index 0000000..759769d
--- /dev/null
+++ b/lib/moped/session/context.rb
@@ -0,0 +1,105 @@
+module Moped
+  class Session
+
+    # @api private
+    class Context
+      extend Forwardable
+
+      def initialize(session)
+        @session = session
+      end
+
+      delegate :safety => :@session
+      delegate :safe?  => :@session
+      delegate :consistency  => :@session
+      delegate :cluster  => :@session
+
+      def login(database, username, password)
+        cluster.auth[database.to_s] = [username, password]
+      end
+
+      def logout(database)
+        cluster.auth.delete database.to_s
+      end
+
+      def query(database, collection, selector, options = {})
+        if consistency == :eventual
+          options[:flags] ||= []
+          options[:flags] |= [:slave_ok]
+        end
+
+        with_node do |node|
+          node.query(database, collection, selector, options)
+        end
+      end
+
+      def command(database, command)
+        options = consistency == :eventual ? { :flags => [:slave_ok] } : {}
+        with_node do |node|
+          node.command(database, command, options)
+        end
+      end
+
+      def insert(database, collection, documents)
+        with_node do |node|
+          if safe?
+            node.pipeline do
+              node.insert(database, collection, documents)
+              node.command("admin", { getlasterror: 1 }.merge(safety))
+            end
+          else
+            node.insert(database, collection, documents)
+          end
+        end
+      end
+
+      def update(database, collection, selector, change, options = {})
+        with_node do |node|
+          if safe?
+            node.pipeline do
+              node.update(database, collection, selector, change, options)
+              node.command("admin", { getlasterror: 1 }.merge(safety))
+            end
+          else
+            node.update(database, collection, selector, change, options)
+          end
+        end
+      end
+
+      def remove(database, collection, selector, options = {})
+        with_node do |node|
+          if safe?
+            node.pipeline do
+              node.remove(database, collection, selector, options)
+              node.command("admin", { getlasterror: 1 }.merge(safety))
+            end
+          else
+            node.remove(database, collection, selector, options)
+          end
+        end
+      end
+
+      def get_more(*args)
+        raise NotImplementedError, "#get_more cannot be called on Context; it must be called directly on a node"
+      end
+
+      def kill_cursors(*args)
+        raise NotImplementedError, "#kill_cursors cannot be called on Context; it must be called directly on a node"
+      end
+
+      def with_node
+        if consistency == :eventual
+          cluster.with_secondary do |node|
+            yield node
+          end
+        else
+          cluster.with_primary do |node|
+            yield node
+          end
+        end
+      end
+
+    end
+
+  end
+end
diff --git a/lib/moped/socket.rb b/lib/moped/socket.rb
deleted file mode 100644
index abb239a..0000000
--- a/lib/moped/socket.rb
+++ /dev/null
@@ -1,201 +0,0 @@
-require "timeout"
-
-module Moped
-
-  # @api private
-  #
-  # The internal class wrapping a socket connection.
-  class Socket
-
-    # Thread-safe atomic integer.
-    class RequestId
-      def initialize
-        @mutex = Mutex.new
-        @id = 0
-      end
-
-      def next
-        @mutex.synchronize { @id += 1 }
-      end
-    end
-
-    attr_reader :connection
-
-    attr_reader :host
-    attr_reader :port
-
-    def initialize(host, port)
-      @host = host
-      @port = port
-
-      @mutex = Mutex.new
-      @request_id = RequestId.new
-    end
-
-    # @return [true, false] whether the connection was successful
-    # @note The connection timeout is currently just 0.5 seconds, which should
-    #   be sufficient, but may need to be raised or made configurable for
-    #   high-latency situations. That said, if connecting to the remote server
-    #   takes that long, we may not want to use the node any way.
-    def connect
-      return true if connection
-
-      Timeout::timeout 0.5 do
-        @connection = TCPSocket.new(host, port)
-      end
-    rescue Errno::ECONNREFUSED, Timeout::Error
-      return false
-    end
-
-    # @return [true, false] whether this socket connection is alive
-    def alive?
-      if connection
-        return false if connection.closed?
-
-        @mutex.synchronize do
-          if select([connection], nil, nil, 0)
-            !connection.eof? rescue false
-          else
-            true
-          end
-        end
-      else
-        false
-      end
-    end
-
-    # Execute the operations on the connection.
-    def execute(*ops)
-      instrument(ops) do
-        buf = ""
-
-        last = ops.each do |op|
-          op.request_id = @request_id.next
-          op.serialize buf
-        end.last
-
-        if Protocol::Query === last || Protocol::GetMore === last
-          length = nil
-
-          @mutex.synchronize do
-            connection.write buf
-
-            length, = connection.read(4).unpack('l<')
-
-            # Re-use the already allocated buffer used for writing the command.
-            connection.read(length - 4, buf)
-          end
-
-          parse_reply length, buf
-        else
-          @mutex.synchronize do
-            connection.write buf
-          end
-
-          nil
-        end
-      end
-    end
-
-    def parse_reply(length, data)
-      buffer = StringIO.new data
-
-      reply = Protocol::Reply.allocate
-
-      reply.length = length
-
-      reply.request_id,
-        reply.response_to,
-        reply.op_code,
-        reply.flags,
-        reply.cursor_id,
-        reply.offset,
-        reply.count = buffer.read(32).unpack('l4<q<l2<')
-
-      reply.documents = reply.count.times.map do
-        BSON::Document.deserialize(buffer)
-      end
-
-      reply
-    end
-
-    # Executes a simple (one result) query and returns the first document.
-    #
-    # @return [Hash] the first document in a result set.
-    def simple_query(query)
-      query = query.dup
-      query.limit = -1
-
-      execute(query).documents.first
-    end
-
-    # Manually closes the connection
-    def close
-      @mutex.synchronize do
-        connection.close if connection && !connection.closed?
-        @connection = nil
-      end
-    end
-
-    def auth
-      @auth ||= {}
-    end
-
-    def apply_auth(credentials)
-      return if auth == credentials
-      logouts = auth.keys - credentials.keys
-
-      logouts.each do |database|
-        logout database
-      end
-
-      credentials.each do |database, (username, password)|
-        login(database, username, password) unless auth[database] == [username, password]
-      end
-    end
-
-    def login(database, username, password)
-      getnonce = Protocol::Command.new(database, getnonce: 1)
-      result = simple_query getnonce
-
-      raise Errors::OperationFailure.new(getnonce, result) unless result["ok"] == 1
-
-      authenticate = Protocol::Commands::Authenticate.new(database, username, password, result["nonce"])
-      result = simple_query authenticate
-      raise Errors::OperationFailure.new(authenticate, result) unless result["ok"] == 1
-
-      auth[database.to_s] = [username, password]
-    end
-
-    def logout(database)
-      command = Protocol::Command.new(database, logout: 1)
-      result = simple_query command
-      raise Errors::OperationFailure.new(command, result) unless result["ok"] == 1
-      auth.delete(database.to_s)
-    end
-
-    def instrument(ops)
-      instrument_start = (logger = Moped.logger) && logger.debug? && Time.now
-      yield
-    ensure
-      log_operations(logger, ops, Time.now - instrument_start) if instrument_start && !$!
-    end
-
-    def log_operations(logger, ops, duration)
-      prefix  = "  MOPED: #{host}:#{port} "
-      indent  = " "*prefix.length
-      runtime = (" (%.1fms)" % duration)
-
-      if ops.length == 1
-        logger.debug prefix + ops.first.log_inspect + runtime
-      else
-        first, *middle, last = ops
-
-        logger.debug prefix + first.log_inspect
-        middle.each { |m| logger.debug indent + m.log_inspect }
-        logger.debug indent + last.log_inspect + runtime
-      end
-    end
-
-  end
-end
diff --git a/lib/moped/threaded.rb b/lib/moped/threaded.rb
new file mode 100644
index 0000000..9576f38
--- /dev/null
+++ b/lib/moped/threaded.rb
@@ -0,0 +1,32 @@
+module Moped
+
+  # This module contains logic for easy access to objects that have a lifecycle
+  # on the current thread.
+  #
+  # Extracted from Mongoid's +Threaded+ module.
+  #
+  # @api private
+  module Threaded
+    extend self
+
+    # Begin a thread-local stack for +name+.
+    def begin(name)
+      stack(name).push true
+    end
+
+    # @return [Boolean] whether the stack is being executed
+    def executing?(name)
+      !stack(name).empty?
+    end
+
+    # End the thread-local stack for +name+.
+    def end(name)
+      stack(name).pop
+    end
+
+    # @return [Array] a named, thread-local stack.
+    def stack(name)
+      Thread.current["[moped]:#{name}-stack"] ||= []
+    end
+  end
+end
diff --git a/perf/cases.rb b/perf/cases.rb
index 8d81d3c..02c1c9f 100644
--- a/perf/cases.rb
+++ b/perf/cases.rb
@@ -55,23 +55,27 @@ profile "Insert and find one (1000x, 1 thread)" do
 end
 
 profile "Insert and find one (1000x, 2 threads)" do
   2.times.map do
     Thread.new do
-      1000.times do
-        session[:people].insert(name: "John")
-        session[:people].find.one
+      session.new do |session|
+        1000.times do
+          session[:people].insert(name: "John")
+          session[:people].find.one
+        end
       end
     end
   end.each &:join
 end
 
 profile "Insert and find one (1000x, 5 threads)" do
   5.times.map do |i|
     Thread.new do
-      1000.times do
-        session[:people].insert(name: "John")
-        session[:people].find.one
+      session.new do |session|
+        1000.times do
+          session[:people].insert(name: "John")
+          session[:people].find.one
+        end
       end
     end
   end.each &:join
 end
diff --git a/spec/integration/protocol/authentication_spec.rb b/spec/integration/protocol/authentication_spec.rb
deleted file mode 100644
index ed10cc3..0000000
--- a/spec/integration/protocol/authentication_spec.rb
+++ /dev/null
@@ -1,103 +0,0 @@
-require "spec_helper"
-
-describe Moped::Protocol do
-  Protocol = Moped::Protocol
-
-  let(:connection) do
-    TCPSocket.new("localhost", 27017)
-  end
-
-  after do
-    connection.close unless connection.closed?
-  end
-
-  describe "authentication" do
-    context "when nonce is invalid" do
-      let(:auth) do
-        Protocol::Commands::Authenticate.new :admin, "user", "pass", "fakenonce"
-      end
-
-      it "fails" do
-        connection.write auth
-        reply = Protocol::Reply.deserialize(connection).documents[0]
-        reply["ok"].should eq 0.0
-      end
-    end
-
-    context "when nonce is valid but user doesn't exist" do
-      let(:nonce) do
-        command = Protocol::Command.new :admin, getnonce: 1
-        connection.write command
-        Protocol::Reply.deserialize(connection).documents[0]["nonce"]
-      end
-
-      let(:auth) do
-        Protocol::Commands::Authenticate.new :admin, "user", "pass", nonce
-      end
-
-      it "fails" do
-        connection.write auth
-        reply = Protocol::Reply.deserialize(connection).documents[0]
-        reply["ok"].should eq 0.0
-      end
-    end
-
-    context "when nonce is valid but password is wrong" do
-      let(:nonce) do
-        command = Protocol::Command.new "moped-protocol-spec", getnonce: 1
-        connection.write command
-        Protocol::Reply.deserialize(connection).documents[0]["nonce"]
-      end
-
-      let(:auth) do
-        Protocol::Commands::Authenticate.new "moped-protocol-spec",
-          "moped",
-          "pass",
-          nonce
-      end
-
-      before do
-        connection.write Protocol::Insert.new(
-          "moped-protocol-spec",
-          "system.users",
-          [{ user: "moped", pwd: Digest::MD5.hexdigest("moped:mongo:password") }]
-        )
-      end
-
-      it "fails" do
-        connection.write auth
-        reply = Protocol::Reply.deserialize(connection).documents[0]
-        reply["ok"].should eq 0.0
-      end
-    end
-
-    context "when authentication is valid" do
-      let(:nonce) do
-        command = Protocol::Command.new "moped-protocol-spec", getnonce: 1
-        connection.write command
-        Protocol::Reply.deserialize(connection).documents[0]["nonce"]
-      end
-
-      let(:auth) do
-        Protocol::Commands::Authenticate.new "moped-protocol-spec",
-          "moped",
-          "password",
-          nonce
-      end
-
-      before do
-        connection.write Protocol::Insert.new(
-          "moped-protocol-spec",
-          "system.users",
-          [{ user: "moped", pwd: Digest::MD5.hexdigest("moped:mongo:password") }]
-        )
-      end
-
-      it "succeeds" do
-        connection.write auth
-        reply = Protocol::Reply.deserialize(connection).documents[0]
-        reply["ok"].should eq 1.0
-      end
-    end
-  end
-end
diff --git a/spec/integration/protocol/protocol_spec.rb b/spec/integration/protocol/protocol_spec.rb
deleted file mode 100644
index 1be93b7..0000000
--- a/spec/integration/protocol/protocol_spec.rb
+++ /dev/null
@@ -1,58 +0,0 @@
-require "spec_helper"
-
-describe Moped::Protocol do
-  let(:Protocol) { Moped::Protocol }
-
-  let(:connection) do
-    TCPSocket.new("localhost", 27017)
-  end
-
-  after do
-    connection.close unless connection.closed?
-  end
-
-  describe "reply response flags" do
-    let(:reply) do
-      Protocol::Reply.deserialize(connection)
-    end
-
-    context "when get more is called with an invalid cursor" do
-      let(:get_more) do
-        Protocol::GetMore.new("moped-protocol", "suite", 0, 0)
-      end
-
-      specify "the cursor not found flag is set" do
-        connection.write get_more
-        reply.flags.should include :cursor_not_found
-      end
-    end
-
-    context "when query generates an error" do
-      let(:query) do
-        Protocol::Query.new "moped-protocol", "people", { '$in' => 1 }, limit: -1
-      end
-
-      specify "the query failure flag is set" do
-        connection.write query
-        reply.flags.should eq [:query_failure]
-      end
-    end
-
-    context "when mongod supports await data query option" do
-      let(:query) do
-        Protocol::Query.new "admin", "$cmd", { buildinfo: 1 }, limit: -1
-      end
-
-      specify "the await capable flag is set" do
-        connection.write query
-
-        if reply.documents[0]["version"] >= "1.6"
-          reply.flags.should eq [:await_capable]
-        else
-          reply.flags.should eq []
-        end
-      end
-    end
-
-  end
-end
diff --git a/spec/integration_spec.rb b/spec/integration_spec.rb
deleted file mode 100644
index 8abc6ec..0000000
--- a/spec/integration_spec.rb
+++ /dev/null
@@ -1,176 +0,0 @@
-# encoding: utf-8
-
-require "spec_helper"
-
-describe Moped::Session do
-  context "with a single master node" do
-    let(:session) { Moped::Session.new ["127.0.0.1:27017"], database: "moped_test" }
-
-    after do
-      session[:people].drop if session[:people].find.count > 0
-      session.cluster.servers.each(&:close)
-    end
-
-    it "inserts and queries a single document" do
-      id = Moped::BSON::ObjectId.new
-      session[:people].insert(_id: id, name: "John")
-      john = session[:people].find(_id: id).one
-      john["_id"].should eq id
-      john["name"].should eq "John"
-    end
-
-    it "inserts and queries on utf-8 data" do
-      id = Moped::BSON::ObjectId.new
-      doc = {
-        "_id" => id,
-        "gültig" => "1",
-        "1" => "gültig",
-        "2" => :"gültig",
-        "3" => ["gültig"],
-        "4" => /gültig/
-      }
-      session[:people].insert(doc)
-      session[:people].find(_id: id).one.should eq doc
-    end
-
-    it "can explain a query" do
-      id = Moped::BSON::ObjectId.new
-      session[:people].find(_id: id).explain["cursor"].should eq("BasicCursor")
-    end
-
-    it "can explain a query with a sort" do
-      id = Moped::BSON::ObjectId.new
-      query = session[:people].find(_id: id)
-      query.sort(_id: 1).explain["cursor"].should eq("BasicCursor")
-    end
-
-    it "drops a collection" do
-      session.command(count: :people)["n"].should eq 0
-      session[:people].insert(name: "John")
-      session.command(count: :people)["n"].should eq 1
-      session[:people].drop
-      session.command(count: :people)["n"].should eq 0
-    end
-
-    it "can be inserted into safely" do
-      session.with(safe: true) do |session|
-        session[:people].insert(name: "John")["ok"].should eq 1
-      end
-    end
-
-    it "raises an error on a failed insert in safe mode" do
-      session.with(safe: true) do |session|
-        lambda do
-          session[:people].insert("$invalid" => nil)
-        end.should raise_exception(Moped::Errors::OperationFailure)
-      end
-    end
-
-    it "can sort documents" do
-      session[:people].insert([{name: "John"}, {name: "Mary"}])
-      session[:people].find.sort(_id: -1).first["name"].should eq "Mary"
-    end
-
-    it "can update documents" do
-      id = Moped::BSON::ObjectId.new
-      session[:people].insert(_id: id, name: "John")
-      mary = session[:people].find(_id: id).one
-      mary["name"].should eq "John"
-      session[:people].find(_id: id).update(name: "Mary")
-      mary = session[:people].find(_id: id).one
-      mary["_id"].should eq id
-      mary["name"].should eq "Mary"
-    end
-
-    it "can update documents safely" do
-      id = Moped::BSON::ObjectId.new
-      session[:people].insert(_id: id, name: "John")
-      mary = session[:people].find(_id: id).one
-      mary["name"].should eq "John"
-      session.with(safe: true) do |session|
-        session[:people].find(_id: id).update(name: "Mary")["ok"].should eq 1
-      end
-    end
-
-    it "can update multiple documents" do
-      session[:people].insert([{name: "John"}, {name: "Mary"}])
-      session[:people].find.update_all("$set" => { "last_name" => "Unknown" })
-
-      session[:people].find.sort(_id: -1).first["last_name"].should eq "Unknown"
-      session[:people].find.sort(_id: 1).first["last_name"].should eq "Unknown"
-    end
-
-    it "can upsert documents" do
-      session[:people].find.upsert(name: "Mary")
-      mary = session[:people].find(name: "Mary").one
-      mary["name"].should eq "Mary"
-    end
-
-    it "can delete a single document" do
-      session[:people].insert([{name: "John"}, {name: "John"}])
-      session[:people].find(name: "John").remove
-      session[:people].find.count.should eq 1
-    end
-
-    it "can delete a single document safely" do
-      session[:people].insert([{name: "John"}, {name: "John"}])
-      session.with(safe: true) do |session|
-        session[:people].find(name: "John").remove["ok"].should eq 1
-      end
-      session[:people].find.count.should eq 1
-    end
-
-    it "can delete a multiple documents" do
-      session[:people].insert([{name: "John"}, {name: "John"}])
-      session[:people].find(name: "John").remove_all
-      session[:people].find.count.should eq 0
-    end
-
-    it "can retrieve multiple documents with fixed limit" do
-      session[:people].insert([{name: "John"}, {name: "Mary"}])
-      john, mary = session[:people].find.limit(-2).sort(name: 1).to_a
-      john["name"].should eq "John"
-      mary["name"].should eq "Mary"
-    end
-
-    it "can retrieve distinct values" do
-      session[:people].insert([{name: "John"}, {name: "Mary"}])
-      values = session[:people].find.distinct(:name)
-      values.should eq [ "John", "Mary" ]
-    end
-
-    it "can retrieve no documents" do
-      session[:people].find.limit(-2).sort(name: 1).to_a.should eq []
-    end
-
-    it "can limit a result set" do
-      documents = 100.times.map { { _id: Moped::BSON::ObjectId.new } }
-      session[:people].insert(documents)
-      session[:people].find.limit(20).to_a.length.should eq 20
-    end
-
-    it "does not leave open cursors" do
-      documents = 100.times.map { { _id: Moped::BSON::ObjectId.new } }
-      session[:people].insert(documents)
-      session[:people].find.limit(20).to_a.length.should eq 20
-      status = session.command serverStatus: 1
-      status["cursors"]["totalOpen"].should eq 0
-    end
-
-    it "can retrieve large result sets" do
-      documents = 1000.times.map do
-        { _id: Moped::BSON::ObjectId.new }
-      end
-      session[:people].insert(documents)
-      session[:people].find.to_a.length.should eq 1000
-    end
-
-    it "can have multiple connections" do
-      status = session.command serverStatus: 1
-      count = status["connections"]["current"]
-      new_session = session.new
-      status = new_session.command serverStatus: 1
-      status["connections"]["current"].should eq count + 1
-    end
-  end
-end
diff --git a/spec/moped/bson/object_id_spec.rb b/spec/moped/bson/object_id_spec.rb
index 5769b00..afb56e6 100644
--- a/spec/moped/bson/object_id_spec.rb
+++ b/spec/moped/bson/object_id_spec.rb
@@ -1,18 +1,18 @@
 require "spec_helper"
 
 describe Moped::BSON::ObjectId do
   let(:bytes) do
-    [78, 77, 102, 52, 59, 57, 182, 132, 7, 0, 0, 1]
+    [78, 77, 102, 52, 59, 57, 182, 132, 7, 0, 0, 1].pack("C12")
   end
 
   describe ".from_string" do
 
     context "when the string is valid" do
 
-      it "initializes with the strings bytes" do
-        Moped::BSON::ObjectId.should_receive(:new).with(bytes)
+      it "initializes with the string's bytes" do
+        Moped::BSON::ObjectId.should_receive(:from_data).with(bytes)
         Moped::BSON::ObjectId.from_string "4e4d66343b39b68407000001"
       end
     end
 
     context "when the string is not valid" do
@@ -45,43 +45,35 @@ describe Moped::BSON::ObjectId do
       end
     end
 
   end
 
-  describe "#initialize" do
-
-    context "with data" do
-      it "sets the object id's data" do
-        Moped::BSON::ObjectId.new(bytes).data.should == bytes
-      end
+  describe "#from_time" do
+    it "sets the generation time" do
+      time = Time.at((Time.now.utc - 64800).to_i).utc
+      Moped::BSON::ObjectId.from_time(time).generation_time.should == time
     end
+  end
 
+  describe "#initialize" do
     context "with no data" do
       it "increments the id on each call" do
         Moped::BSON::ObjectId.new.should_not eq Moped::BSON::ObjectId.new
       end
 
       it "increments the id safely across threads" do
         ids = 2.times.map { Thread.new { Moped::BSON::ObjectId.new } }
         ids[0].value.should_not eq ids[1].value
       end
     end
-
-    context "with a time" do
-      it "sets the generation time" do
-        time = Time.at((Time.now.utc - 64800).to_i).utc
-        Moped::BSON::ObjectId.new(nil, time).generation_time.should == time
-      end
-    end
-
   end
 
   describe "#==" do
 
     context "when data is identical" do
       it "returns true" do
-        Moped::BSON::ObjectId.new(bytes).should == Moped::BSON::ObjectId.new(bytes)
+        Moped::BSON::ObjectId.from_data(bytes).should == Moped::BSON::ObjectId.from_data(bytes)
       end
     end
 
     context "when other is not an object id" do
       it "returns false" do
@@ -93,11 +85,11 @@ describe Moped::BSON::ObjectId do
 
   describe "#eql?" do
 
     context "when data is identical" do
       it "returns true" do
-        Moped::BSON::ObjectId.new(bytes).should eql Moped::BSON::ObjectId.new(bytes)
+        Moped::BSON::ObjectId.from_data(bytes).should eql Moped::BSON::ObjectId.from_data(bytes)
       end
     end
 
     context "when other is not an object id" do
       it "returns false" do
@@ -109,11 +101,11 @@ describe Moped::BSON::ObjectId do
 
   describe "#hash" do
 
     context "when data is identical" do
       it "returns the same hash" do
-        Moped::BSON::ObjectId.new(bytes).hash.should eq Moped::BSON::ObjectId.new(bytes).hash
+        Moped::BSON::ObjectId.from_data(bytes).hash.should eq Moped::BSON::ObjectId.from_data(bytes).hash
       end
     end
 
     context "when other is not an object id" do
       it "returns a different hash" do
@@ -124,11 +116,11 @@ describe Moped::BSON::ObjectId do
   end
 
   describe "#to_s" do
 
     it "returns a hex string representation of the id" do
-      Moped::BSON::ObjectId.new(bytes).to_s.should eq "4e4d66343b39b68407000001"
+      Moped::BSON::ObjectId.from_data(bytes).to_s.should eq "4e4d66343b39b68407000001"
     end
 
   end
 
 end
diff --git a/spec/moped/cluster_spec.rb b/spec/moped/cluster_spec.rb
index 4ee42a7..37eda5f 100644
--- a/spec/moped/cluster_spec.rb
+++ b/spec/moped/cluster_spec.rb
@@ -1,358 +1,319 @@
 require "spec_helper"
 
-describe Moped::Cluster do
-
-  let(:master) do
-    TCPServer.new "127.0.0.1", 0
-  end
-
-  let(:secondary_1) do
-    TCPServer.new "127.0.0.1", 0
-  end
-
-  let(:secondary_2) do
-    TCPServer.new "127.0.0.1", 0
-  end
-
-  describe "initialize" do
-    let(:cluster) do
-      Moped::Cluster.new(["127.0.0.1:27017","127.0.0.1:27018"], true)
-    end
-
-    it "stores the list of seeds" do
-      cluster.seeds.should eq ["127.0.0.1:27017", "127.0.0.1:27018"]
-    end
-
-    it "stores whether the connection is direct" do
-      cluster.direct.should be_true
-    end
-
-    it "has an empty list of primaries" do
-      cluster.primaries.should be_empty
-    end
-
-    it "has an empty list of secondaries" do
-      cluster.secondaries.should be_empty
-    end
-
-    it "has an empty list of servers" do
-      cluster.servers.should be_empty
-    end
-
-    it "has an empty list of dynamic seeds" do
-      cluster.dynamic_seeds.should be_empty
-    end
+describe Moped::Cluster, replica_set: true do
+  let(:replica_set) do
+    Moped::Cluster.new(seeds, {})
   end
 
-  describe "#sync" do
-    let(:cluster) { Moped::Cluster.new(["127.0.0.1:27017"]) }
-
-    it "syncs each seed node" do
-      server = Moped::Server.allocate
-      Moped::Server.should_receive(:new).with("127.0.0.1:27017").and_return(server)
-
-      cluster.should_receive(:sync_server).with(server).and_return([])
-      cluster.sync
-    end
-  end
-
-  describe "#sync_server" do
-    let(:cluster) { Moped::Cluster.new [""], false }
-    let(:server) { Moped::Server.new("localhost:27017") }
-    let(:socket) { Moped::Socket.new "", 99999 }
-    let(:connection) { Support::MockConnection.new }
-
-    before do
-      socket.stub(connection: connection, alive?: true)
-      server.stub(socket: socket)
+  context "when the replica set hasn't connected yet" do
+    describe "#with_primary" do
+      it "connects and yields the primary node" do
+        replica_set.with_primary do |node|
+          node.address.should eq @primary.address
+        end
+      end
     end
 
-    context "when node is not running" do
-      it "returns nothing" do
-        socket.stub(connect: false)
-
-        cluster.sync_server(server).should be_empty
+    describe "#with_secondary" do
+      it "connects and yields a secondary node" do
+        replica_set.with_secondary do |node|
+          @secondaries.map(&:address).should include node.address
+        end
       end
     end
 
-    context "when talking to a single node" do
+    context "and the primary is down" do
       before do
-        connection.pending_replies << Hash[
-          "ismaster" => true,
-          "maxBsonObjectSize" => 16777216,
-          "ok" => 1.0
-        ]
-      end
-
-      it "adds the node to the master set" do
-        cluster.sync_server server
-        cluster.primaries.should include server
+        @primary.stop
       end
-    end
-
-    context "when talking to a replica set node" do
-
-      context "that is not configured" do
-        before do
-          connection.pending_replies << Hash[
-            "ismaster" => false,
-            "secondary" => false,
-            "info" => "can't get local.system.replset config from self or any seed (EMPTYCONFIG)",
-            "isreplicaset" => true,
-            "maxBsonObjectSize" => 16777216,
-            "ok" => 1.0
-          ]
-        end
 
-        it "returns nothing" do
-          cluster.sync_server(server).should be_empty
+      describe "#with_primary" do
+        it "raises a connection error" do
+          lambda do
+            replica_set.with_primary do |node|
+              node.command "admin", ping: 1
+            end
+          end.should raise_exception(Moped::Errors::ConnectionFailure)
         end
       end
 
-      context "that is being initiated" do
-        before do
-          connection.pending_replies << Hash[
-            "ismaster" => false,
-            "secondary" => false,
-            "info" => "Received replSetInitiate - should come online shortly.",
-            "isreplicaset" => true,
-            "maxBsonObjectSize" => 16777216,
-            "ok" => 1.0
-          ]
-        end
-
-        it "raises a connection failure exception" do
-          cluster.sync_server(server).should be_empty
+      describe "#with_secondary" do
+        it "connects and yields a secondary node" do
+          replica_set.with_secondary do |node|
+            @secondaries.map(&:address).should include node.address
+          end
         end
       end
+    end
 
-      context "that is ready but not elected" do
-        before do
-          connection.pending_replies << Hash[
-            "setName" => "3fef4842b608",
-            "ismaster" => false,
-            "secondary" => false,
-            "hosts" => ["localhost:61085", "localhost:61086", "localhost:61084"],
-            "primary" => "localhost:61084",
-            "me" => "localhost:61085",
-            "maxBsonObjectSize" => 16777216,
-            "ok" => 1.0
-          ]
-        end
-
-        it "raises no exception" do
-          lambda do
-            cluster.sync_server server
-          end.should_not raise_exception
-        end
+    context "and a single secondary is down" do
+      before do
+        @secondaries.first.stop
+      end
 
-        it "adds the server to the list" do
-          cluster.sync_server server
-          cluster.servers.should include server
+      describe "#with_primary" do
+        it "connects and yields the primary node" do
+          replica_set.with_primary do |node|
+            node.address.should eq @primary.address
+          end
         end
+      end
 
-        it "returns all other known hosts" do
-          cluster.sync_server(server).should =~ ["localhost:61085", "localhost:61086", "localhost:61084"]
+      describe "#with_secondary" do
+        it "connects and yields a secondary node" do
+          replica_set.with_secondary do |node|
+            node.address.should eq @secondaries.last.address
+          end
         end
       end
+    end
 
-      context "that is ready" do
-        before do
-          connection.pending_replies << Hash[
-            "setName" => "3ff029114780",
-            "ismaster" => true,
-            "secondary" => false,
-            "hosts" => ["localhost:59246", "localhost:59248", "localhost:59247"],
-            "primary" => "localhost:59246",
-            "me" => "localhost:59246",
-            "maxBsonObjectSize" => 16777216,
-            "ok" => 1.0
-          ]
-        end
+    context "and all secondaries are down" do
+      before do
+        @secondaries.each &:stop
+      end
 
-        it "adds the node to the master set" do
-          cluster.sync_server server
-          cluster.primaries.should include server
+      describe "#with_primary" do
+        it "connects and yields the primary node" do
+          replica_set.with_primary do |node|
+            node.address.should eq @primary.address
+          end
         end
+      end
 
-        it "returns all other known hosts" do
-          cluster.sync_server(server).should =~ ["localhost:59246", "localhost:59248", "localhost:59247"]
+      describe "#with_secondary" do
+        it "connects and yields the primary node" do
+          replica_set.with_secondary do |node|
+            node.address.should eq @primary.address
+          end
         end
       end
-
     end
   end
 
-  describe "#socket_for" do
-    let(:cluster) do
-      Moped::Cluster.new ""
+  context "when the replica set is connected" do
+    before do
+      replica_set.refresh
     end
 
-    let(:server) do
-      Moped::Server.new("localhost:27017").tap do |server|
-        server.stub(socket: socket)
+    describe "#with_primary" do
+      it "connects and yields the primary node" do
+        replica_set.with_primary do |node|
+          node.address.should eq @primary.address
+        end
       end
     end
 
-    let(:socket) do
-      Moped::Socket.new("127.0.0.1", 27017).tap do |socket|
-        socket.connect
+    describe "#with_secondary" do
+      it "connects and yields a secondary node" do
+        replica_set.with_secondary do |node|
+          @secondaries.map(&:address).should include node.address
+        end
       end
     end
 
-    context "when socket is dead" do
-      let(:dead_server) do
-        Moped::Server.allocate.tap do |server|
-          server.stub(socket: dead_socket)
+    context "and the primary is down" do
+      before do
+        @primary.stop
+      end
+
+      describe "#with_primary" do
+        it "raises a connection error" do
+          lambda do
+            replica_set.with_primary do |node|
+              node.command "admin", ping: 1
+            end
+          end.should raise_exception(Moped::Errors::ConnectionFailure)
         end
       end
 
-      let(:dead_socket) do
-        Moped::Socket.new("127.0.0.1", 27017).tap do |socket|
-          socket.stub(:alive? => false)
+      describe "#with_secondary" do
+        it "connects and yields a secondary node" do
+          replica_set.with_secondary do |node|
+            @secondaries.map(&:address).should include node.address
+          end
         end
       end
+    end
 
+    context "and a single secondary is down" do
       before do
-        primaries = [server, dead_server]
-        primaries.stub(:sample).and_return(dead_server, server)
-        cluster.stub(primaries: primaries)
+        @secondaries.first.stop
       end
 
-      it "removes the socket" do
-        cluster.should_receive(:remove).with(dead_server)
-        cluster.socket_for :write
+      describe "#with_primary" do
+        it "connects and yields the primary node" do
+          replica_set.with_primary do |node|
+            node.address.should eq @primary.address
+          end
+        end
       end
 
-      it "returns the living socket" do
-        cluster.socket_for(:write).should eq socket
+      describe "#with_secondary" do
+        it "connects and yields a secondary node" do
+          replica_set.with_secondary do |node|
+            node.command "admin", ping: 1
+            node.address.should eq @secondaries.last.address
+          end
+        end
       end
     end
 
-    context "when mode is write" do
+    context "and all secondaries are down" do
       before do
-        server.primary = true
+        @secondaries.each &:stop
       end
 
-      context "and the cluster is not synced" do
-        it "syncs the cluster" do
-          cluster.should_receive(:sync) do
-            cluster.servers << server
+      describe "#with_primary" do
+        it "connects and yields the primary node" do
+          replica_set.with_primary do |node|
+            node.address.should eq @primary.address
           end
-          cluster.socket_for :write
         end
+      end
 
-        it "returns the socket" do
-          cluster.stub(:sync) { cluster.servers << server }
-          cluster.socket_for(:write).should eq socket
+      describe "#with_secondary" do
+        it "connects and yields the primary node" do
+          replica_set.with_secondary do |node|
+            node.command "admin", ping: 1
+            node.address.should eq @primary.address
+          end
         end
+      end
+    end
+  end
 
-        it "applies the cached authentication" do
-          cluster.stub(:sync) { cluster.servers << server }
-          socket.should_receive(:apply_auth).with(cluster.auth)
-          cluster.socket_for(:write)
-        end
+  context "with down interval" do
+    let(:replica_set) do
+      Moped::Cluster.new(seeds, { down_interval: 5 })
+    end
+
+    context "and all secondaries are down" do
+      before do
+        replica_set.refresh
+        @secondaries.each &:stop
+        replica_set.refresh
       end
 
-      context "and the cluster is synced" do
-        before do
-          cluster.servers << server
+      describe "#with_secondary" do
+        it "connects and yields the primary node" do
+          replica_set.with_secondary do |node|
+            node.command "admin", ping: 1
+            node.address.should eq @primary.address
+          end
         end
+      end
 
-        it "does not re-sync the cluster" do
-          cluster.should_receive(:sync).never
-          cluster.socket_for :write
+      context "when a secondary node comes back up" do
+        before do
+          @secondaries.each &:restart
         end
 
-        it "returns the socket" do
-          cluster.socket_for(:write).should eq socket
+        describe "#with_secondary" do
+          it "connects and yields the primary node" do
+            replica_set.with_secondary do |node|
+              node.command "admin", ping: 1
+              node.address.should eq @primary.address
+            end
+          end
         end
 
-        it "applies the cached authentication" do
-          socket.should_receive(:apply_auth).with(cluster.auth)
-          cluster.socket_for(:write)
+        context "and the node is ready to be retried" do
+          it "connects and yields the secondary node" do
+            Time.stub(:new).and_return(Time.now + 10)
+            replica_set.with_secondary do |node|
+              node.command "admin", ping: 1
+              @secondaries.map(&:address).should include node.address
+            end
+          end
         end
       end
     end
+  end
 
-    context "when mode is read" do
-      context "and the cluster is not synced" do
-        before do
-          server.primary = true
-        end
+  context "with only primary provided as a seed" do
+    let(:replica_set) do
+      Moped::Cluster.new([@primary.address], {})
+    end
 
-        it "syncs the cluster" do
-          cluster.should_receive(:sync) do
-            cluster.servers << server
-          end
-          cluster.socket_for :read
+    describe "#with_primary" do
+      it "connects and yields the primary node" do
+        replica_set.with_primary do |node|
+          node.address.should eq @primary.address
         end
+      end
+    end
 
-        it "applies the cached authentication" do
-          cluster.stub(:sync) { cluster.servers << server }
-          socket.should_receive(:apply_auth).with(cluster.auth)
-          cluster.socket_for(:read)
+    describe "#with_secondary" do
+      it "connects and yields a secondary node" do
+        replica_set.with_secondary do |node|
+          @secondaries.map(&:address).should include node.address
         end
       end
+    end
+  end
 
-      context "and the cluster is synced" do
-        context "and no secondaries are found" do
-          before do
-            server.primary = true
-            cluster.servers << server
-          end
-
-          it "returns the master connection" do
-            cluster.socket_for(:read).should eq socket
-          end
+  context "with only primary provided as a seed" do
+    let(:replica_set) do
+      Moped::Cluster.new([@secondaries[0].address], {})
+    end
 
-          it "applies the cached authentication" do
-            socket.should_receive(:apply_auth).with(cluster.auth)
-            cluster.socket_for(:read)
-          end
+    describe "#with_primary" do
+      it "connects and yields the primary node" do
+        replica_set.with_primary do |node|
+          node.address.should eq @primary.address
         end
+      end
+    end
 
-        context "and a slave is found" do
-          it "returns a random slave connection" do
-            secondaries = [server]
-            cluster.stub(secondaries: secondaries)
-            secondaries.should_receive(:sample).and_return(server)
-            cluster.socket_for(:read).should eq socket
-          end
-
-          it "applies the cached authentication" do
-            cluster.stub(secondaries: [server])
-            socket.should_receive(:apply_auth).with(cluster.auth)
-            cluster.socket_for(:read)
-          end
+    describe "#with_secondary" do
+      it "connects and yields a secondary node" do
+        replica_set.with_secondary do |node|
+          @secondaries.map(&:address).should include node.address
         end
       end
     end
   end
+end
 
-  describe "#login" do
-    let(:cluster) do
-      Moped::Cluster.allocate
-    end
+describe Moped::Cluster, "authentication", mongohq: :auth do
+  let(:session) do
+    Support::MongoHQ.auth_session(false)
+  end
 
-    it "adds the credentials to the auth cache" do
-      cluster.login("admin", "username", "password")
-      cluster.auth.should eq("admin" => ["username", "password"])
+  describe "logging in with valid credentials" do
+    it "logs in and processes commands" do
+      session.login *Support::MongoHQ.auth_credentials
+      session.command(ping: 1).should eq("ok" => 1)
     end
   end
 
-  describe "#logout" do
-    let(:cluster) do
-      Moped::Cluster.allocate
+  describe "logging in with invalid credentials" do
+    it "raises an AuthenticationFailure exception" do
+      session.login "invalid-user", "invalid-password"
+
+      lambda do
+        session.command(ping: 1)
+      end.should raise_exception(Moped::Errors::AuthenticationFailure)
     end
+  end
 
+  describe "logging in with valid credentials and then logging out" do
     before do
-      cluster.login("admin", "username", "password")
+      session.login *Support::MongoHQ.auth_credentials
+      session.command(ping: 1).should eq("ok" => 1)
     end
 
-    it "removes the stored credentials" do
-      cluster.logout :admin
-      cluster.auth.should be_empty
+    it "logs out" do
+      lambda do
+        session.command dbStats: 1
+      end.should_not raise_exception
+
+      session.logout
+
+      lambda do
+        session.command dbStats: 1
+      end.should raise_exception(Moped::Errors::OperationFailure)
     end
   end
 end
diff --git a/spec/moped/collection_spec.rb b/spec/moped/collection_spec.rb
index b3db63f..e995b96 100644
--- a/spec/moped/collection_spec.rb
+++ b/spec/moped/collection_spec.rb
@@ -1,92 +1,39 @@
 require "spec_helper"
 
 describe Moped::Collection do
-
   let(:session) do
-    mock(Moped::Session)
-  end
-
-  let(:database) do
-    mock(Moped::Database, session: session, name: "moped")
-  end
-
-  let(:collection) do
-    described_class.new database, :users
+    Moped::Session.new %w[127.0.0.1:27017], database: "moped_test"
   end
 
-  describe "#initialize" do
-
-    it "stores the database" do
-      collection.database.should eq database
-    end
-
-    it "stores the collection name" do
-      collection.name.should eq :users
-    end
-  end
-
-  describe "#indexes" do
-    it "returns a new indexes instance" do
-      collection.indexes.should be_an_instance_of Moped::Indexes
-    end
-  end
+  let(:scope) { object_id }
 
   describe "#drop" do
-
-    it "drops the collection" do
-      database.should_receive(:command).with(drop: :users)
-      collection.drop
-    end
-  end
-
-  describe "#find" do
-
-    let(:selector) do
-      Hash[ a: 1 ]
-    end
-
-    let(:query) do
-      mock(Moped::Query)
-    end
-
-    it "returns a new Query" do
-      Moped::Query.should_receive(:new).
-        with(collection, selector).and_return(query)
-      collection.find(selector).should eq query
+    before do
+      session.drop
+      session.command create: "users"
     end
 
-    it "defaults to an empty selector" do
-      Moped::Query.should_receive(:new).
-        with(collection, {}).and_return(query)
-      collection.find.should eq query
+    it "drops the collection" do
+      result = session[:users].drop
+      result["ns"].should eq "moped_test.users"
     end
   end
 
   describe "#insert" do
-
-    before do
-      session.should_receive(:with, :consistency => :strong).and_yield(session)
-      session.stub safe?: false
-    end
-
-    context "when passed a single document" do
-
-      it "inserts the document" do
-        session.should_receive(:execute).with do |insert|
-          insert.documents.should eq [{a: 1}]
-        end
-        collection.insert(a: 1)
-      end
+    it "inserts a single document" do
+      document = { "_id" => Moped::BSON::ObjectId.new, "scope" => scope }
+      session[:users].insert(document)
+      session[:users].find(document).one.should eq document
     end
 
-    context "when passed multiple documents" do
+    it "insert multiple documents" do
+      documents = [
+        { "_id" => Moped::BSON::ObjectId.new, "scope" => scope },
+        { "_id" => Moped::BSON::ObjectId.new, "scope" => scope }
+      ]
 
-      it "inserts the documents" do
-        session.should_receive(:execute).with do |insert|
-          insert.documents.should eq [{a: 1}, {b: 2}]
-        end
-        collection.insert([{a: 1}, {b: 2}])
-      end
+      session[:users].insert(documents)
+      session[:users].find(scope: scope).entries.should eq documents
     end
   end
 end
diff --git a/spec/moped/cursor_spec.rb b/spec/moped/cursor_spec.rb
deleted file mode 100644
index e4f2aa9..0000000
--- a/spec/moped/cursor_spec.rb
+++ /dev/null
@@ -1,237 +0,0 @@
-require "spec_helper"
-
-describe Moped::Cursor do
-  let(:session) { mock Moped::Session }
-  let(:query_operation) { Moped::Protocol::Query.allocate }
-  let(:cursor) { Moped::Cursor.new(session, query_operation) }
-
-  describe "#initialize" do
-    it "stores the session" do
-      cursor.session.should eq session
-    end
-
-    it "stores a copy of the query operation" do
-      query_operation.should_receive(:dup).and_return(query_operation)
-      cursor.query_op.should eq query_operation
-    end
-
-    describe "the get_more operation" do
-      it "inherits the query's database" do
-        cursor.get_more_op.database.should eq query_operation.database
-      end
-
-      it "inherits the query's collection" do
-        cursor.get_more_op.collection.should eq query_operation.collection
-      end
-
-      it "inherits the query's limit" do
-        cursor.get_more_op.limit.should eq query_operation.limit
-      end
-    end
-  end
-
-  describe "#more?" do
-    context "when get more operation's cursor id is 0" do
-      it "returns false" do
-        cursor.get_more_op.cursor_id = 0
-        cursor.more?.should be_false
-      end
-    end
-    context "when get more operation's cursor id is not 0" do
-      it "returns true" do
-        cursor.get_more_op.cursor_id = 123
-        cursor.more?.should be_true
-      end
-    end
-  end
-
-  describe "#limited?" do
-    context "when original query's limit is greater than 0" do
-      before do
-        query_operation.limit = 20
-      end
-
-      it "returns true" do
-        cursor.should be_limited
-      end
-    end
-
-    context "when original query's limit is not greater than 0" do
-      before do
-        query_operation.limit = 0
-      end
-
-      it "returns true" do
-        cursor.should_not be_limited
-      end
-    end
-  end
-
-  describe "#query" do
-    let(:reply) do
-      Moped::Protocol::Reply.allocate.tap do |reply|
-        reply.cursor_id = 123
-        reply.count = 1
-        reply.documents = [{"a" => 1}]
-      end
-    end
-
-    before do
-      session.stub(query: reply)
-    end
-
-    context "when query is limited" do
-      before do
-        query_operation.limit = 21
-        cursor.query query_operation
-      end
-
-      it "updates the more operation's limit" do
-        cursor.get_more_op.limit.should eq 20
-      end
-
-      it "sets the kill cursor operation's cursor id" do
-        cursor.kill_cursor_op.cursor_ids.should eq [reply.cursor_id]
-      end
-
-      it "sets the more operation's cursor id" do
-        cursor.get_more_op.cursor_id.should eq reply.cursor_id
-      end
-    end
-
-    context "when query is limited" do
-      before do
-        query_operation.limit = 0
-        cursor.query query_operation
-      end
-
-      it "does not update the more operation's limit" do
-        cursor.get_more_op.limit.should eq query_operation.limit
-      end
-
-      it "sets the kill cursor operation's cursor id" do
-        cursor.kill_cursor_op.cursor_ids.should eq [reply.cursor_id]
-      end
-
-      it "sets the more operation's cursor id" do
-        cursor.get_more_op.cursor_id.should eq reply.cursor_id
-      end
-    end
-
-    it "returns the documents" do
-      cursor.query(query_operation).should eq reply.documents
-    end
-  end
-
-  describe "#each" do
-
-    context "when query returns all available documents" do
-      let(:reply) do
-        Moped::Protocol::Reply.allocate.tap do |reply|
-          reply.cursor_id = 0
-          reply.count = 21
-          reply.documents = [{"a" => 1}]
-        end
-      end
-
-      before do
-        session.stub(query: reply)
-      end
-
-      it "yields each document" do
-        results = []
-        cursor.each { |doc| results << doc }
-        results.should eq reply.documents
-      end
-
-      it "does not get more" do
-        session.should_receive(:query).once
-        cursor.each {}
-      end
-
-      it "does not kill the cursor" do
-        cursor.should_receive(:kill).never
-        cursor.each {}
-      end
-    end
-
-    context "when query is unlimited" do
-      let(:reply) do
-        Moped::Protocol::Reply.allocate.tap do |reply|
-          reply.cursor_id = 10
-          reply.count = 10
-          reply.documents = [{"a" => 1}]
-        end
-      end
-
-      let(:get_more_reply) do
-        Moped::Protocol::Reply.allocate.tap do |reply|
-          reply.cursor_id = 0
-          reply.count = 21
-          reply.documents = [{"a" => 1}]
-        end
-      end
-
-      before do
-        session.stub(:query).and_return(reply, get_more_reply)
-      end
-
-      it "yields each document" do
-        results = []
-        cursor.each { |doc| results << doc }
-        results.should eq reply.documents + get_more_reply.documents
-      end
-
-      it "gets more twice" do
-        session.should_receive(:query).twice
-        cursor.each {}
-      end
-
-      it "does not kill the cursor" do
-        cursor.should_receive(:kill).never
-        cursor.each {}
-      end
-    end
-
-    context "when query is limited" do
-      let(:reply) do
-        Moped::Protocol::Reply.allocate.tap do |reply|
-          reply.cursor_id = 10
-          reply.count = 10
-          reply.documents = [{"a" => 1}]
-        end
-      end
-
-      let(:get_more_reply) do
-        Moped::Protocol::Reply.allocate.tap do |reply|
-          reply.cursor_id = 10
-          reply.count = 10
-          reply.documents = [{"a" => 1}]
-        end
-      end
-
-      before do
-        query_operation.limit = 20
-        session.stub(:query).and_return(reply, get_more_reply)
-        session.stub(:execute)
-      end
-
-      it "yields each document" do
-        results = []
-        cursor.each { |doc| results << doc }
-        results.should eq reply.documents + get_more_reply.documents
-      end
-
-      it "gets more twice" do
-        session.should_receive(:query).at_least(2)
-        cursor.each {}
-      end
-
-      it "kills the cursor" do
-        cursor.should_receive(:kill).once
-        cursor.each {}
-      end
-    end
-
-  end
-end
diff --git a/spec/moped/database_spec.rb b/spec/moped/database_spec.rb
deleted file mode 100644
index df0ea83..0000000
--- a/spec/moped/database_spec.rb
+++ /dev/null
@@ -1,86 +0,0 @@
-require "spec_helper"
-
-describe Moped::Database do
-
-  let(:session) do
-    Moped::Session.new ""
-  end
-
-  let(:database) do
-    Moped::Database.new(session, :admin)
-  end
-
-  describe "#initialize" do
-
-    it "stores the session" do
-      database.session.should eq session
-    end
-
-    it "stores the database name" do
-      database.name.should eq :admin
-    end
-  end
-
-  describe "#command" do
-
-    before do
-      session.stub(:with).and_yield(session)
-    end
-
-    it "runs the given command against the master connection" do
-      session.should_receive(:with, :consistency => :strong).
-        and_yield(session)
-      session.should_receive(:simple_query) do |query|
-        query.full_collection_name.should eq "admin.$cmd"
-        query.selector.should eq(ismaster: 1)
-
-        { "ok" => 1.0 }
-      end
-
-      database.command ismaster: 1
-    end
-
-    context "when the command fails" do
-
-      it "raises an exception" do
-        session.stub(simple_query: { "ok" => 0.0 })
-
-        expect {
-          database.command ismaster: 1
-        }.to raise_exception(Moped::Errors::OperationFailure)
-      end
-    end
-  end
-
-  describe "#drop" do
-
-    it "drops the database" do
-      database.should_receive(:command).with(dropDatabase: 1)
-      database.drop
-    end
-  end
-
-  describe "#[]" do
-
-    it "returns a collection with that name" do
-      Moped::Collection.should_receive(:new).with(database, :users)
-      database[:users]
-    end
-  end
-
-  describe "#login" do
-
-    it "logs in to the database with the username and password" do
-      session.cluster.should_receive(:login).with(:admin, "username", "password")
-      database.login("username", "password")
-    end
-  end
-
-  describe "#log out" do
-
-    it "logs out from the database" do
-      session.cluster.should_receive(:logout).with(:admin)
-      database.logout
-    end
-  end
-end
diff --git a/spec/moped/errors_spec.rb b/spec/moped/errors_spec.rb
deleted file mode 100644
index 932c9da..0000000
--- a/spec/moped/errors_spec.rb
+++ /dev/null
@@ -1,91 +0,0 @@
-require "spec_helper"
-
-describe Moped::Errors do
-
-  describe "OperationFailure" do
-    let(:command) do
-      Moped::Protocol::Query.allocate
-    end
-
-    let(:error_details) do
-      { "$err"=>"invalid query", "code"=>12580 }
-    end
-
-    let(:error) do
-      described_class::OperationFailure.new(command, error_details)
-    end
-
-    describe "#initialize" do
-      it "stores the command which generated the error" do
-        error.command.should eq command
-      end
-
-      it "stores the details about the error" do
-        error.details.should eq error_details
-      end
-    end
-
-    describe "#message" do
-      it "includes the command that generated the error" do
-        error.message.should include command.inspect
-      end
-
-      context "when code is included in error details" do
-        let(:error_details) do
-          { "err" => "invalid query", "code" => 12580 }
-        end
-
-        it "includes the code" do
-          error.message.should include error_details["code"].to_s
-        end
-
-        it "includes the error code reference site" do
-          error.message.should include Moped::Errors::ERROR_REFERENCE
-        end
-
-        it "includes the error message" do
-          error.message.should include error_details["err"].inspect
-        end
-      end
-
-      context "when err is in the error details" do
-        let(:error_details) do
-          { "err" => "invalid query" }
-        end
-
-        it "includes the error message" do
-          error.message.should include error_details["err"].inspect
-        end
-      end
-
-      context "when $err is in the error details" do
-        let(:error_details) do
-          { "$err" => "not master" }
-        end
-
-        it "includes the error message" do
-          error.message.should include error_details["$err"].inspect
-        end
-      end
-
-      context "when errmsg is in the error details" do
-        let(:error_details) do
-          { "errmsg" => "invalid query" }
-        end
-
-        it "includes the error message" do
-          error.message.should include error_details["errmsg"].inspect
-        end
-      end
-    end
-
-  end
-
-  describe "QueryFailure" do
-    it "is a kind of OperationFailure" do
-      Moped::Errors::QueryFailure.ancestors.should \
-        include Moped::Errors::OperationFailure
-    end
-  end
-
-end
diff --git a/spec/moped/indexes_spec.rb b/spec/moped/indexes_spec.rb
index cfc643a..50cf353 100644
--- a/spec/moped/indexes_spec.rb
+++ b/spec/moped/indexes_spec.rb
@@ -1,128 +1,53 @@
 require "spec_helper"
 
 describe Moped::Indexes do
-  let(:session) { Moped::Session.new ["127.0.0.1:27017"], database: "moped_test" }
-  let(:indexes) do
-    described_class.new(session.current_database, :users)
-  end
-
-  after do
-    session.command(deleteIndexes: "users", index: "*")
+  let(:session) do
+    Moped::Session.new %w[127.0.0.1:27017], database: "moped_test"
   end
 
-  describe "#each" do
-    before do
-      session[:"system.indexes"].insert(ns: "moped_test.users", key: { name: 1 }, name: "name_1")
-    end
-
-    it "yields all indexes on the collection" do
-      indexes.to_a.should eq \
-        session[:"system.indexes"].find(ns: "moped_test.users").to_a
-    end
+  let(:indexes) do
+    session[:users].indexes
   end
 
-  describe "#[]" do
-    before do
-      session[:"system.indexes"].insert(ns: "moped_test.users", key: { name: 1 }, name: "name_1")
-    end
-
-    it "returns the index with the provided key" do
-      indexes[name: 1]["name"].should eq "name_1"
-    end
+  before do
+    indexes.drop
   end
 
   describe "#create" do
-    let(:key) do
-      Hash["location.latlong" => "2d", "name" => 1, "age" => -1]
-    end
-
-    context "with no options" do
-      it "creates an index with a generated name" do
-        indexes.create(key)
-        indexes[key]["name"].should eq "location.latlong_2d_name_1_age_-1"
-      end
-    end
-
-    context "with a name provided" do
-      it "creates an index with the provided name" do
-        indexes.create(key, name: "custom_index_name")
-        indexes[key]["name"].should eq "custom_index_name"
-      end
-    end
-
-    context "with background: true" do
-      it "creates an index" do
-        indexes.create(key, background: true)
-        indexes[key]["background"].should eq true
-      end
-    end
-
-    context "with dropDups: true" do
-      it "creates an index" do
-        indexes.create(key, dropDups: true)
-        indexes[key]["dropDups"].should eq true
+    context "when called without extra options" do
+      it "creates an index with no options" do
+        indexes.create name: 1
+        indexes[name: 1].should_not be_nil
       end
     end
 
-    context "with unique: true" do
-      it "creates an index" do
-        indexes.create(key, unique: true)
-        indexes[key]["unique"].should eq true
+    context "when called with extra options" do
+      it "creates an index with the extra options" do
+        indexes.create({name: 1}, {unique: true, dropDups: true})
+        index = indexes[name: 1]
+        index["unique"].should be_true
+        index["dropDups"].should be_true
       end
     end
-
-    context "with sparse: true" do
-      it "creates an index" do
-        indexes.create(key, sparse: true)
-        indexes[key]["sparse"].should eq true
-      end
-    end
-
-    context "with v: 0" do
-      it "creates an index" do
-        indexes.create(key, v: 0)
-        indexes[key]["v"].should eq 0
-      end
-    end
-
   end
 
   describe "#drop" do
-    before do
-      indexes.create name: 1
-      indexes.create age: -1
-    end
-
-    context "with no key" do
-      before do
-        indexes.drop
-      end
-
-      it "drops all indexes for the collection" do
-        indexes[name: 1].should be_nil
-        indexes[age: -1].should be_nil
+    context "when provided a key" do
+      it "drops the index" do
+        indexes.create name: 1
+        indexes.drop(name: 1).should be_true
       end
     end
 
-    context "with a key" do
-      before do
-        indexes.drop(name: 1)
-      end
-
-      it "drops the index that matches the key" do
+    context "when not provided a key" do
+      it "drops all indexes" do
+        indexes.create name: 1
+        indexes.create age: 1
+        indexes.drop
         indexes[name: 1].should be_nil
-      end
-
-      it "does not drop other indexes" do
-        indexes[age: -1].should_not be_nil
-      end
-    end
-
-    context "with a key that doesn't exist" do
-      it "returns false" do
-        indexes.drop(other: 1).should be_false
+        indexes[age: 1].should be_nil
       end
     end
   end
 
 end
diff --git a/spec/moped/logging_spec.rb b/spec/moped/logging_spec.rb
deleted file mode 100644
index 9b82f79..0000000
--- a/spec/moped/logging_spec.rb
+++ /dev/null
@@ -1,83 +0,0 @@
-require "spec_helper"
-
-describe Moped::Logging do
-  let(:config) do
-    Module.new { extend Moped::Logging }
-  end
-  let(:logger) { mock(Logger) }
-
-  describe ".rails_logger" do
-    context "when Rails is present" do
-      let(:rails) { Class.new }
-
-      before do
-        Object.const_set :Rails, rails
-      end
-
-      after do
-        Object.send(:remove_const, :Rails)
-      end
-
-      context "and it defines logger" do
-        before do
-          rails.stub(logger: logger)
-        end
-
-        it "returns the logger" do
-          config.rails_logger.should eq logger
-        end
-      end
-
-      context "but does not define logger" do
-        it "returns false" do
-          config.rails_logger.should be_false
-        end
-      end
-
-    end
-  end
-
-  describe ".default_logger" do
-    it "returns a new logger instance" do
-      config.default_logger.should be_a_kind_of Logger
-    end
-
-    it "sets the log level to info" do
-      config.default_logger.level.should eq Logger::INFO
-    end
-  end
-
-  describe ".logger" do
-    context "when a rails logger is available" do
-      before do
-        config.stub(rails_logger: logger)
-      end
-
-      it "returns the rails logger" do
-        config.logger.should eq logger
-      end
-    end
-
-    context "when a rails logger is not available" do
-      before do
-        config.stub(rails_logger: nil)
-        config.stub(default_logger: logger)
-      end
-
-      it "returns the default logger" do
-        config.logger.should eq logger
-      end
-    end
-
-    context "when the logger is set to nil" do
-      before do
-        config.logger = nil
-      end
-
-      it "returns nil" do
-        config.logger.should be_nil
-      end
-    end
-  end
-
-end
diff --git a/spec/moped/node_spec.rb b/spec/moped/node_spec.rb
new file mode 100644
index 0000000..2a0dfcf
--- /dev/null
+++ b/spec/moped/node_spec.rb
@@ -0,0 +1,57 @@
+require "spec_helper"
+
+describe Moped::Node, replica_set: true do
+  let(:replica_set_node) do
+    @replica_set.nodes.first
+  end
+
+  let(:node) do
+    Moped::Node.new(replica_set_node.address)
+  end
+
+  describe "#ensure_connected" do
+    context "when node is running" do
+      it "processes the block" do
+        node.ensure_connected do
+          node.command("admin", ping: 1)
+        end.should eq("ok" => 1)
+      end
+    end
+
+    context "when node is not running" do
+      before do
+        replica_set_node.stop
+      end
+
+      it "raises a connection error" do
+        lambda do
+          node.ensure_connected do
+            node.command("admin", ping: 1)
+          end
+        end.should raise_exception(Moped::Errors::ConnectionFailure)
+      end
+
+      it "marks the node as down" do
+        node.ensure_connected {} rescue nil
+        node.should be_down
+      end
+    end
+
+    context "when node is connected but connection is dropped" do
+      before do
+        node.ensure_connected do
+          node.command("admin", ping: 1)
+        end
+
+        replica_set_node.hiccup
+      end
+
+      it "reconnects without raising an exception" do
+        node.ensure_connected do
+          node.command("admin", ping: 1)
+        end.should eq("ok" => 1)
+      end
+    end
+  end
+
+end
diff --git a/spec/moped/query_spec.rb b/spec/moped/query_spec.rb
index f3ecb6f..aff94b8 100644
--- a/spec/moped/query_spec.rb
+++ b/spec/moped/query_spec.rb
@@ -1,288 +1,365 @@
 require "spec_helper"
 
 describe Moped::Query do
+  shared_examples_for "Query" do
+    let(:scope) do
+      object_id
+    end
 
-  let(:session) do
-    mock(Moped::Session)
-  end
-
-  let(:database) do
-    mock(
-      Moped::Database,
-      name: "moped",
-      session: session
-    )
-  end
-
-  let(:collection) do
-    mock(
-      Moped::Collection,
-      database: database,
-      name: "users"
-    )
-  end
-
-  let(:selector) do
-    Hash[ a: 1 ]
-  end
-
-  let(:query) do
-    described_class.new collection, selector
-  end
-
-  describe "#initialize" do
-
-    it "stores the collection" do
-      query.collection.should eq collection
+    before do
+      users.find.remove_all
     end
 
-    it "stores the selector" do
-      query.selector.should eq selector
+    let(:documents) do
+      [
+        { "_id" => Moped::BSON::ObjectId.new, "scope" => scope },
+        { "_id" => Moped::BSON::ObjectId.new, "scope" => scope }
+      ]
     end
-  end
 
-  describe "#limit" do
+    it "raises a query failure exception for invalid queries" do
+      lambda do
+        users.find("age" => { "$in" => nil }).first
+      end.should raise_exception(Moped::Errors::QueryFailure)
+    end
 
-    it "sets the query operation's limit field" do
-      query.limit(5)
-      query.operation.limit.should eq 5
+    describe "#limit" do
+      it "limits the query" do
+        users.insert(documents)
+        users.find(scope: scope).limit(1).to_a.should eq [documents.first]
+      end
     end
 
-    it "returns the query" do
-      query.limit(5).should eql query
+    describe "#skip" do
+      it "skips +n+ documents" do
+        users.insert(documents)
+        users.find(scope: scope).skip(1).to_a.should eq [documents.last]
+      end
     end
-  end
 
-  describe "#skip" do
+    describe "#sort" do
+      let(:documents) do
+        [
+          { "_id" => Moped::BSON::ObjectId.new, "scope" => scope, "n" => 0 },
+          { "_id" => Moped::BSON::ObjectId.new, "scope" => scope, "n" => 1 }
+        ]
+      end
 
-    it "sets the query operation's skip field" do
-      query.skip(5)
-      query.operation.skip.should eq 5
+      it "sorts the results" do
+        users.insert(documents)
+        users.find(scope: scope).sort(n: -1).to_a.should eq documents.reverse
+      end
     end
 
-    it "returns the query" do
-      query.skip(5).should eql query
+    describe "#distinct" do
+      let(:documents) do
+        [
+          { count: 0, scope: scope },
+          { count: 1, scope: scope },
+          { count: 1, scope: scope }
+        ]
+      end
+
+      it "returns distinct values for +key+" do
+        users.insert(documents)
+        users.find(scope: scope).distinct(:count).should =~ [0, 1]
+      end
     end
-  end
 
-  describe "#select" do
+    describe "#select" do
+      let(:documents) do
+        [
+          { "scope" => scope, "n" => 0 },
+          { "scope" => scope, "n" => 1 }
+        ]
+      end
 
-    it "sets the query operation's fields" do
-      query.select(a: 1)
-      query.operation.fields.should eq(a: 1)
+      it "changes the fields returned" do
+        users.insert(documents)
+        users.find(scope: scope).select(_id: 0).to_a.should eq documents
+      end
     end
 
-    it "returns the query" do
-      query.select(a: 1).should eql query
-    end
-  end
+    describe "#one" do
+      before do
+        users.insert(documents)
+      end
 
-  describe "#sort" do
+      it "returns the first matching document" do
+        users.find(scope: scope).one.should eq documents.first
+      end
 
-    context "when called for the first time" do
+      it "respects #skip" do
+        users.find(scope: scope).skip(1).one.should eq documents.last
+      end
 
-      it "updates the selector to mongo's advanced selector" do
-        query.sort(a: 1)
-        query.operation.selector.should eq(
-          "$query" => selector,
-          "$orderby" => { a: 1 }
-        )
+      it "respects #sort" do
+        users.find(scope: scope).sort(_id: -1).one.should eq documents.last
       end
     end
 
-    context "when called again" do
+    describe "#explain" do
+      context "when a sort exists" do
+        it "updates to a mongo advanced selector" do
+          stats = Support::Stats.collect do
+            users.find(scope: scope).sort(_id: 1).explain
+          end
+
+          operation = stats[node_for_reads].grep(Moped::Protocol::Query).last
+          operation.selector.should eq(
+            "$query" => { scope: scope },
+            "$explain" => true,
+            "$orderby" => { _id: 1 }
+          )
+        end
+      end
 
-      it "changes the $orderby" do
-        query.sort(a: 1)
-        query.sort(a: 2)
-        query.operation.selector.should eq(
-          "$query" => selector,
-          "$orderby" => { a: 2 }
-        )
+      context "when no sort exists" do
+        it "updates to a mongo advanced selector" do
+          stats = Support::Stats.collect do
+            users.find(scope: scope).explain
+          end
+
+          operation = stats[node_for_reads].grep(Moped::Protocol::Query).last
+          operation.selector.should eq(
+            "$query" => { scope: scope },
+            "$explain" => true,
+            "$orderby" => {}
+          )
+        end
       end
     end
 
-    it "returns the query" do
-      query.sort(a: 1).should eql query
-    end
-  end
+    describe "#each" do
+      it "yields each document" do
+        users.insert(documents)
+        users.find(scope: scope).each.with_index do |document, index|
+          document.should eq documents[index]
+        end
+      end
 
-  describe "#explain" do
+      context "with a limit" do
+        it "closes open cursors" do
+          users.insert(100.times.map { Hash["scope" => scope] })
 
-    before do
-      session.should_receive(:simple_query).with(query.operation)
-    end
+          stats = Support::Stats.collect do
+            users.find(scope: scope).limit(5).entries
+          end
 
-    context "when a sort exists" do
+          stats[node_for_reads].grep(Moped::Protocol::KillCursors).count.should eq 1
+        end
 
-      before do
-        query.sort(_id: 1)
       end
 
-      it "updates to a mongo advanced selector" do
-        query.explain
-        query.operation.selector.should eq(
-          "$query" => selector,
-          "$explain" => true,
-          "$orderby" => { _id: 1 }
-        )
-      end
-    end
+      context "without a limit" do
+        it "fetches more" do
+          users.insert(102.times.map { Hash["scope" => scope] })
 
-    context "when no sort exists" do
+          stats = Support::Stats.collect do
+            users.find(scope: scope).entries
+          end
 
-      it "updates to a mongo advanced selector" do
-        query.explain
-        query.operation.selector.should eq(
-          "$query" => selector,
-          "$explain" => true,
-          "$orderby" => {}
-        )
+          stats[node_for_reads].grep(Moped::Protocol::GetMore).count.should eq 1
+        end
       end
     end
-  end
 
-  describe "#one" do
+    describe "#count" do
+      let(:documents) do
+        [
+          { "_id" => Moped::BSON::ObjectId.new, "scope" => scope },
+          { "_id" => Moped::BSON::ObjectId.new, "scope" => scope },
+          { "_id" => Moped::BSON::ObjectId.new }
+        ]
+      end
 
-    it "executes a simple query" do
-      session.should_receive(:simple_query).with(query.operation)
-      query.one
+      it "returns the number of matching document" do
+        users.insert(documents)
+        users.find(scope: scope).count.should eq 2
+      end
     end
-  end
 
-  describe "#distinct" do
+    describe "#update" do
+      it "updates the first matching document" do
+        users.insert(documents)
+        users.find(scope: scope).update("$set" => { "updated" => true })
+        users.find(scope: scope, updated: true).count.should eq 1
+      end
+    end
 
-    it "executes a distinct command" do
-      database.should_receive(:command).with(
-        distinct: collection.name,
-        key: "name",
-        query: selector
-      ).and_return("values" => [ "durran", "bernerd" ])
-      query.distinct(:name)
+    describe "#update_all" do
+      it "updates all matching documents" do
+        users.insert(documents)
+        users.find(scope: scope).update_all("$set" => { "updated" => true })
+        users.find(scope: scope, updated: true).count.should eq 2
+      end
     end
-  end
 
-  describe "#count" do
+    describe "#upsert" do
+      context "when a document exists" do
+        before do
+          users.insert(scope: scope, counter: 1)
+        end
 
-    it "executes a count command" do
-      database.should_receive(:command).with(
-        count: collection.name,
-        query: selector
-      ).and_return("n" => 4)
+        it "updates the document" do
+          users.find(scope: scope).upsert("$inc" => { counter: 1 })
+          users.find(scope: scope).one["counter"].should eq 2
+        end
+      end
 
-      query.count
+      context "when no document exists" do
+        it "inserts a document" do
+          users.find(scope: scope).upsert("$inc" => { counter: 1 })
+          users.find(scope: scope).one["counter"].should eq 1
+        end
+      end
     end
 
-    it "returns the count" do
-      database.stub(command: { "n" => 4 })
+    describe "#remove" do
+      it "removes the first matching document" do
+        users.insert(documents)
+        users.find(scope: scope).remove
+        users.find(scope: scope).count.should eq 1
+      end
+    end
 
-      query.count.should eq 4
+    describe "#remove_all" do
+      it "removes all matching documents" do
+        users.insert(documents)
+        users.find(scope: scope).remove_all
+        users.find(scope: scope).count.should eq 0
+      end
     end
   end
 
-  describe "#update" do
-
-    let(:change) do
-      Hash[ a: 1 ]
+  context "with a local connection" do
+    let(:session) do
+      Moped::Session.new %w[127.0.0.1:27017], database: "moped_test"
     end
 
-    it "updates the record matching selector with change" do
-      session.should_receive(:with, :consistency => :strong).
-        and_yield(session)
+    let(:users) { session[:users] }
+    let(:node_for_reads) { :primary }
 
-      session.should_receive(:execute).with do |update|
-        update.flags.should eq []
-        update.selector.should eq query.operation.selector
-        update.update.should eq change
-      end
-
-      query.update change
-    end
-  end
+    include_examples "Query"
 
-  describe "#update_all" do
+    describe "#each" do
+      context "with a limit and large result set" do
+        it "gets more and closes cursors" do
+          11.times do
+            users.insert(scope: scope, large_field: "a"*1_000_000)
+          end
 
-    let(:change) do
-      Hash[ a: 1 ]
-    end
+          stats = Support::Stats.collect do
+            users.find(scope: scope).limit(10).entries
+          end
 
-    it "updates all records matching selector with change" do
-      query.should_receive(:update).with(change, [:multi])
-      query.update_all change
+          stats[:primary].grep(Moped::Protocol::GetMore).count.should eq 1
+          stats[:primary].grep(Moped::Protocol::KillCursors).count.should eq 1
+        end
+      end
     end
   end
 
-  describe "#upsert" do
-
-    let(:change) do
-      Hash[ a: 1 ]
+  context "with a remote connection", mongohq: :auth do
+    before :all do
+      @session = Support::MongoHQ.auth_session
     end
 
-    it "upserts the record matching selector with change" do
-      query.should_receive(:update).with(change, [:upsert])
-      query.upsert change
-    end
+    let(:users) { @session[:users] }
+    let(:node_for_reads) { :primary }
+
+    include_examples "Query"
   end
 
-  describe "#remove" do
+  context "with a remote replica set connection with eventual consistency", mongohq: :replica_set do
+    before :all do
+      @session = Support::MongoHQ.replica_set_session.with(safe: true, consistency: :eventual)
+      @session.command ping: 1
+    end
 
-    it "removes the first matching document" do
-      session.should_receive(:with, :consistency => :strong).
-        and_yield(session)
+    let(:users) { @session[:users] }
+    let(:node_for_reads) { :secondary }
 
-      session.should_receive(:execute).with do |delete|
-        delete.flags.should eq [:remove_first]
-        delete.selector.should eq query.operation.selector
-      end
-
-      query.remove
-    end
+    include_examples "Query"
   end
 
-  describe "#remove_all" do
+  context "with a remote replica set connection with strong consistency", mongohq: :replica_set do
+    before :all do
+      @session = Support::MongoHQ.replica_set_session.with(safe: true, consistency: :strong)
+    end
 
-    it "removes all matching documents" do
-      session.should_receive(:with, :consistency => :strong).
-        and_yield(session)
+    let(:users) { @session[:users] }
+    let(:node_for_reads) { :primary }
 
-      session.should_receive(:execute).with do |delete|
-        delete.flags.should eq []
-        delete.selector.should eq query.operation.selector
-      end
+    include_examples "Query"
+  end
 
-      query.remove_all
+  context "with a local replica set w/ failover", replica_set: true do
+    let(:session) do
+      Moped::Session.new seeds, database: "moped_test"
     end
-  end
 
-  describe "#each" do
+    let(:scope) do
+      object_id
+    end
 
     before do
-      session.should_receive(:with).
-        with(retain_socket: true).and_return(session)
+      # Force connection before recording stats
+      session.command ping: 1
     end
 
-    it "creates a new cursor" do
-      cursor = mock(Moped::Cursor, next: nil)
-      Moped::Cursor.should_receive(:new).
-        with(session, query.operation).and_return(cursor)
+    context "and running with eventual consistency" do
+      it "queries a secondary node" do
+        stats = Support::Stats.collect do
+          session.with(consistency: :eventual)[:users].find(scope: scope).entries
+        end
 
-      query.each
-    end
+        stats[:secondary].grep(Moped::Protocol::Query).count.should eq 1
+        stats[:primary].should be_empty
+      end
+
+      it "sets the slave ok flag" do
+        stats = Support::Stats.collect do
+          session.with(consistency: :eventual)[:users].find(scope: scope).one
+        end
+
+        query = stats[:secondary].grep(Moped::Protocol::Query).first
+        query.flags.should include :slave_ok
+      end
 
-    it "yields all documents in the cursor" do
-      cursor = Moped::Cursor.allocate
-      cursor.stub(:to_enum).and_return([1, 2].to_enum)
+      context "and no secondaries are available" do
+        before do
+          @secondaries.each &:stop
+        end
 
-      Moped::Cursor.stub(new: cursor)
+        it "queries the primary node" do
+          stats = Support::Stats.collect do
+            session.with(consistency: :eventual)[:users].find(scope: scope).entries
+          end
 
-      query.to_a.should eq [1, 2]
+          stats[:primary].grep(Moped::Protocol::Query).count.should eq 1
+        end
+      end
     end
 
-    it "returns an enumerator" do
-      cursor = mock(Moped::Cursor)
-      Moped::Cursor.stub(new: cursor)
+    context "and running with strong consistency" do
+      it "queries the primary node" do
+        stats = Support::Stats.collect do
+          session.with(consistency: :strong)[:users].find(scope: scope).entries
+        end
+
+        stats[:primary].grep(Moped::Protocol::Query).count.should eq 1
+        stats[:secondary].should be_empty
+      end
 
-      query.each.should be_a Enumerator
+      it "does not set the slave ok flag" do
+        stats = Support::Stats.collect do
+          session.with(consistency: :strong)[:users].find(scope: scope).one
+        end
+
+        query = stats[:primary].grep(Moped::Protocol::Query).first
+        query.flags.should_not include :slave_ok
+      end
     end
   end
 end
diff --git a/spec/moped/server_spec.rb b/spec/moped/server_spec.rb
deleted file mode 100644
index c9f652b..0000000
--- a/spec/moped/server_spec.rb
+++ /dev/null
@@ -1,80 +0,0 @@
-require "spec_helper"
-
-describe Moped::Server do
-
-  describe "#initialize" do
-    let(:server) do
-      described_class.new("localhost:123")
-    end
-
-    it "stores the original address" do
-      server.address.should eq "localhost:123"
-    end
-
-    it "stores the resolved address" do
-      server.resolved_address.should eql "127.0.0.1:123"
-    end
-
-    it "stores the resolved ip" do
-      server.ip_address.should eq "127.0.0.1"
-    end
-
-    it "stores the port" do
-      server.port.should eq 123
-    end
-  end
-
-  describe "==" do
-    context "when ip and port are the same" do
-      it "returns true" do
-        described_class.new("127.0.0.1:999").should eq \
-          described_class.new("localhost:999")
-      end
-    end
-
-    context "when ip and port are different" do
-      it "returns false" do
-        described_class.new("127.0.0.1:1000").should_not eq \
-          described_class.new("localhost:999")
-      end
-    end
-
-    context "when other is not a server" do
-      it "returns false" do
-        described_class.new("127.0.0.1:999").should_not eq 1
-      end
-    end
-  end
-
-  context "when added to a set" do
-    let(:set) { Set.new }
-
-    context "and ip and port are the same" do
-      it "does not add both servers" do
-        set << described_class.new("127.0.0.1:1000")
-        set << described_class.new("127.0.0.1:1000")
-
-        set.length.should eq 1
-      end
-
-      context "and the original address is different" do
-        it "does not add both servers" do
-          set << described_class.new("localhost:1000")
-          set << described_class.new("127.0.0.1:1000")
-
-          set.length.should eq 1
-        end
-      end
-    end
-
-    context "and ip and port are different" do
-      it "adds both servers" do
-        set << described_class.new("127.0.0.1:1000")
-        set << described_class.new("127.0.0.1:2000")
-
-        set.length.should eq 2
-      end
-    end
-  end
-
-end
diff --git a/spec/moped/session_spec.rb b/spec/moped/session_spec.rb
index c389dd9..e164203 100644
--- a/spec/moped/session_spec.rb
+++ b/spec/moped/session_spec.rb
@@ -1,548 +1,63 @@
 require "spec_helper"
 
 describe Moped::Session do
-
-  let(:seeds) do
-    "127.0.0.1:27017"
-  end
-
-  let(:options) do
-    Hash[database: "test", safe: true, consistency: :eventual]
-  end
-
   let(:session) do
-    described_class.new seeds, options
-  end
-
-  describe "#initialize" do
-
-    it "stores the options provided" do
-      session.options.should eq(options)
-    end
-
-    it "stores the cluster" do
-      session.cluster.should be_a(Moped::Cluster)
-    end
-  end
-
-  describe "#current_database" do
-
-    context "when no database option has been set" do
-
-      let(:session) do
-        described_class.new seeds, {}
-      end
-
-      it "raises an exception" do
-        expect { session.current_database }.to raise_exception
-      end
-    end
-
-    context "when a database option is set" do
-
-      let(:database) do
-        stub
-      end
-
-      before do
-        Moped::Database.should_receive(:new).
-          with(session, options[:database]).and_return(database)
-      end
-
-      it "returns the database from the options" do
-        session.current_database.should eq(database)
-      end
-
-      it "memoizes the database" do
-        database = session.current_database
-        session.current_database.should equal(database)
-      end
-    end
-  end
-
-  describe "#safe?" do
-
-    context "when :safe is not present" do
-
-      before do
-        session.options.delete(:safe)
-      end
-
-      it "returns false" do
-        session.should_not be_safe
-      end
-    end
-
-    context "when :safe is present but false" do
-
-      before do
-        session.options[:safe] = false
-      end
-
-      it "returns false" do
-        session.should_not be_safe
-      end
-    end
-
-    context "when :safe is true" do
-
-      before do
-        session.options[:safe] = true
-      end
-
-      it "returns true" do
-        session.should be_safe
-      end
-    end
-
-    context "when :safe is a hash" do
-
-      before do
-        session.options[:safe] = { fsync: true }
-      end
-
-      it "returns true" do
-        session.should be_safe
-      end
-    end
+    Moped::Session.new %w[127.0.0.1:27017], database: "moped_test"
   end
 
   describe "#use" do
-
-    it "sets the :database option" do
-      session.use :admin
-      session.options[:database].should eq(:admin)
-    end
-
-    context "when there is not already a current database" do
-
-      it "sets the current database" do
-        session.should_receive(:set_current_database).with(:admin)
-        session.use :admin
-      end
+    it "changes the current database" do
+      session.use "moped_test_2"
+      session.command(dbStats: 1)["db"].should eq "moped_test_2"
     end
   end
 
   describe "#with" do
-
-    let(:new_options) do
-      Hash[database: "test-2"]
-    end
-
     context "when called with a block" do
-
-      it "yields a session" do
-        session.with(new_options) do |new_session|
-          new_session.should be_a Moped::Session
-        end
+      it "returns the value from the block" do
+        session.with { :value }.should eq :value
       end
 
-      it "yields a new session" do
-        session.with(new_options) do |new_session|
-          new_session.should_not eql session
+      it "yields a session with the provided options" do
+        session.with(safe: true) do |safe|
+          safe.options[:safe].should eq true
         end
       end
 
-      it "returns the result of the block" do
-        session.with(new_options) { false }.should eq false
-      end
-
-      it "merges the old and new session's options" do
-        session.with(new_options) do |new_session|
-          new_session.options.should eq options.merge(new_options)
-        end
-      end
-
-      it "does not change the original session's options" do
-        original_options = options.dup
-        session.with(new_options) do |new_session|
-          session.options.should eql original_options
-        end
-      end
-
-      it "unmemoizes the current database" do
-        db = session.current_database
-        session.with(new_options) do |new_session|
-          new_session.current_database.should_not eql db
+      it "does not modify the original session" do
+        session.with(database: "other") do |safe|
+          session.options[:database].should eq "moped_test"
         end
       end
     end
 
     context "when called without a block" do
-
-      it "returns a session" do
-        session.with(new_options).should be_a Moped::Session
-      end
-
-      it "returns a new session" do
-        session.with(new_options).should_not eql session
-      end
-
-      it "merges the old and new session's options" do
-        session.with(new_options).options.should eq options.merge(new_options)
-      end
-
-      it "does not change the original session's options" do
-        original_options = options.dup
-        session.with(new_options)
-        session.options.should eql original_options
-      end
-    end
-  end
-
-  describe "#new" do
-
-    let(:new_options) do
-      Hash[database: "test-2"]
-    end
-
-    let(:new_session) do
-      described_class.new seeds, options
-    end
-
-    before do
-      new_session.cluster.stub(:reconnect)
-    end
-
-    it "delegates to #with" do
-      session.should_receive(:with).with(new_options).and_return(new_session)
-      session.new(new_options)
-    end
-
-    it "instructs the cluster to reconnect" do
-      session.stub(with: new_session)
-      new_session.cluster.should_receive(:reconnect)
-      session.new(new_options)
-    end
-
-    context "when called with a block" do
-
-      it "yields the new session" do
-        session.stub(with: new_session)
-        session.new(new_options) do |session|
-          session.should eql new_session
-        end
+      it "returns a session with the provided options" do
+        safe = session.with(safe: true)
+        safe.options[:safe].should eq true
       end
-    end
 
-    context "when called without a block" do
-
-      it "returns the new session" do
-        session.stub(with: new_session)
-        session.new(new_options).should eql new_session
+      it "does not modify the original session" do
+        other = session.with(database: "other")
+        session.options[:database].should eq "moped_test"
       end
     end
   end
 
   describe "#drop" do
-
-    it "delegates to the current database" do
-      database = mock(Moped::Database)
-      session.should_receive(:current_database).and_return(database)
-      database.should_receive(:drop)
-      session.drop
-    end
-  end
-
-  describe "#command" do
-
-    let(:command) do
-      Hash[ismaster: 1]
-    end
-
-    it "delegates to the current database" do
-      database = mock(Moped::Database)
-      session.should_receive(:current_database).and_return(database)
-      database.should_receive(:command).with(command)
-      session.command command
-    end
-  end
-
-  describe "#login" do
-
-    it "delegates to the current database" do
-      database = mock(Moped::Database)
-      session.should_receive(:current_database).and_return(database)
-      database.should_receive(:login).with("username", "password")
-      session.login("username", "password")
-    end
-  end
-
-  describe "#logout" do
-
-    it "delegates to the current database" do
-      database = mock(Moped::Database)
-      session.should_receive(:current_database).and_return(database)
-      database.should_receive(:logout)
-      session.logout
-    end
-  end
-
-  describe "#socket_for" do
-
-    it "delegates to the cluster" do
-      session.cluster.should_receive(:socket_for).with(:read)
-      session.send(:socket_for, :read)
-    end
-
-    context "when retain socket option is set" do
-
-      before do
-        session.options[:retain_socket] = true
-      end
-
-      it "only aquires the socket once" do
-        session.cluster.should_receive(:socket_for).
-          with(:read).once.and_return(mock(Moped::Socket))
-
-        session.send(:socket_for, :read)
-        session.send(:socket_for, :read)
-      end
-    end
-  end
-
-  describe "#simple_query" do
-
-    let(:query) do
-      Moped::Protocol::Query.allocate
-    end
-
-    let(:socket) do
-      mock(Moped::Socket)
-    end
-
-    let(:reply) do
-      Moped::Protocol::Reply.allocate.tap do |reply|
-        reply.documents = [{a: 1}]
+    it "drops the current database" do
+      session.with(database: "moped_test_2") do |session|
+        session.drop.should eq("dropped" => "moped_test_2", "ok" => 1)
       end
     end
-
-    before do
-      session.stub(socket_for: socket)
-      session.stub(query: reply)
-    end
-
-    it "limits the query" do
-      session.should_receive(:query) do |query|
-        query.limit.should eq(-1)
-        reply
-      end
-
-      session.simple_query(query)
-    end
-
-    it "returns the document" do
-      session.simple_query(query).should eq(a: 1)
-    end
   end
 
-  describe "#query" do
-
-    let(:query) do
-      Moped::Protocol::Query.allocate
-    end
-
-    let(:socket) do
-      mock(Moped::Socket)
-    end
-
-    let(:reply) do
-      Moped::Protocol::Reply.allocate.tap do |reply|
-        reply.documents = [{a: 1}]
-      end
-    end
-
-    before do
-      session.stub(socket_for: socket)
-      socket.stub(:execute).and_return(reply)
-    end
-
-    context "when consistency is strong" do
-
-      before do
-        session.options[:consistency] = :strong
-      end
-
-      it "queries the master node" do
-        session.should_receive(:socket_for).with(:write).
-          and_return(socket)
-        session.query(query)
-      end
-    end
-
-    context "when consistency is eventual" do
-
-      before do
-        session.options[:consistency] = :eventual
-      end
-
-      it "queries a slave node" do
-        session.should_receive(:socket_for).with(:read).
-          and_return(socket)
-        session.query(query)
-      end
-
-      context "and query accepts flags" do
-
-        it "sets slave_ok on the query flags" do
-          session.stub(socket_for: socket)
-          socket.should_receive(:execute) do |query|
-            query.flags.should include :slave_ok
-          end
-
-          session.query(query)
-        end
-      end
-
-      context "and query does not accept flags" do
-
-        let(:query) do
-          Moped::Protocol::GetMore.allocate
-        end
-
-        it "doesn't try to set flags" do
-          session.stub(socket_for: socket)
-          expect { session.query(query) }.not_to raise_exception
-        end
-      end
-    end
-
-    context "when reply has :query_failure flag" do
-
-      before do
-        reply.flags = [:query_failure]
-      end
-
-      it "raises a QueryFailure exception" do
-        expect {
-          session.query(query)
-        }.to raise_exception(Moped::Errors::QueryFailure)
+  describe "#command" do
+    it "runs the command on the current database" do
+      session.with(database: "moped_test_2") do |session|
+        session.command(dbStats: 1)["db"].should eq "moped_test_2"
       end
     end
   end
 
-  describe "#execute" do
-
-    let(:operation) do
-      Moped::Protocol::Insert.allocate
-    end
-
-    let(:socket) do
-      mock(Moped::Socket)
-    end
-
-    context "when session is not in safe mode" do
-
-      before do
-        session.options[:safe] = false
-      end
-
-      context "when consistency is strong" do
-
-        before do
-          session.options[:consistency] = :strong
-        end
-
-        it "executes the operation on the master node" do
-          session.should_receive(:socket_for).with(:write).
-            and_return(socket)
-          socket.should_receive(:execute).with(operation)
-
-          session.execute(operation)
-        end
-      end
-
-      context "when consistency is eventual" do
-
-        before do
-          session.options[:consistency] = :eventual
-        end
-
-        it "executes the operation on a slave node" do
-          session.should_receive(:socket_for).with(:read).
-            and_return(socket)
-          socket.should_receive(:execute).with(operation)
-
-          session.execute(operation)
-        end
-      end
-    end
-
-    context "when session is in safe mode" do
-
-      let(:reply) do
-        Moped::Protocol::Reply.allocate.tap do |reply|
-          reply.documents = [{a: 1}]
-        end
-      end
-
-      before do
-        session.options[:safe] = { w: 2 }
-      end
-
-      context "when the operation fails" do
-
-        let(:reply) do
-          Moped::Protocol::Reply.allocate.tap do |reply|
-            reply.documents = [{
-              "err"=>"document to insert can't have $ fields",
-              "code"=>13511,
-              "n"=>0,
-              "connectionId"=>894,
-              "ok"=>1.0
-            }]
-          end
-        end
-
-        it "raises an OperationFailure exception" do
-          session.stub(socket_for: socket)
-          socket.stub(execute: reply)
-
-          expect {
-            session.execute(operation)
-          }.to raise_exception(Moped::Errors::OperationFailure)
-        end
-      end
-
-      context "when consistency is strong" do
-
-        before do
-          session.options[:consistency] = :strong
-        end
-
-        it "executes the operation on the master node" do
-          session.should_receive(:socket_for).with(:write).
-            and_return(socket)
-
-          socket.should_receive(:execute) do |op, query|
-            op.should eq operation
-            query.selector.should eq(getlasterror: 1, w: 2)
-            reply
-          end
-
-          session.execute(operation)
-        end
-      end
-
-      context "when consistency is eventual" do
-
-        before do
-          session.options[:consistency] = :eventual
-        end
-
-        it "executes the operation on a slave node" do
-          session.should_receive(:socket_for).with(:read).
-            and_return(socket)
-
-          socket.should_receive(:execute) do |op, query|
-            op.should eq operation
-            query.selector.should eq(getlasterror: 1, w: 2)
-            reply
-          end
-
-          session.execute(operation)
-        end
-      end
-    end
-  end
 end
diff --git a/spec/moped/socket_spec.rb b/spec/moped/socket_spec.rb
deleted file mode 100644
index 22be808..0000000
--- a/spec/moped/socket_spec.rb
+++ /dev/null
@@ -1,462 +0,0 @@
-require "spec_helper"
-
-describe Moped::Socket do
-
-  let!(:server) do
-    TCPServer.new "127.0.0.1", 0
-  end
-
-  let(:socket) do
-    described_class.new "127.0.0.1", server.addr[1]
-  end
-
-  let(:connection) {
-    socket.connection
-  }
-
-  before do
-    socket.connect
-  end
-
-  after do
-    connection.close if connection && !connection.closed?
-    server.close unless server.closed?
-  end
-
-  describe "#initialize" do
-    it "stores the host of the server" do
-      socket.host.should eq "127.0.0.1"
-    end
-
-    it "stores the port of the server" do
-      socket.port.should eq server.addr[1]
-    end
-
-    it "connects to the server" do
-      socket.connection.should_not be_closed
-    end
-  end
-
-  describe "#connect" do
-    context "when node is not running" do
-      let(:bogus_port) do
-        server = TCPServer.new("127.0.0.1", 0)
-        server.addr[1].tap do
-          server.close
-        end
-      end
-
-      let(:socket) do
-        described_class.new "127.0.0.1", bogus_port
-      end
-
-      it "returns false" do
-        socket.connect.should be_false
-      end
-    end
-
-    context "when connection times out" do
-      if RUBY_PLATFORM == "java"
-        let(:timeout_server) do
-          java.net.ServerSocket.new(0, 1)
-        end
-
-        let(:timeout_port) do
-          timeout_server.getLocalPort
-        end
-      else
-        let(:timeout_server) do
-          TCPServer.new("127.0.0.1", 0).tap do |server|
-            server.listen(1)
-          end
-        end
-
-        let(:timeout_port) do
-          timeout_server.addr[1]
-        end
-      end
-
-      let(:timeout_socket) do
-        described_class.new "127.0.0.1", timeout_port
-      end
-
-      before do
-        sockaddr = Socket.pack_sockaddr_in(timeout_port, '127.0.0.1')
-
-        5.times do # flood the server socket
-          ::Socket.new(::Socket::AF_INET, ::Socket::SOCK_STREAM, 0).connect_nonblock(sockaddr) rescue nil
-        end
-      end
-
-      after do
-        timeout_server.close unless timeout_server.closed?
-      end
-
-      it "returns false" do
-        timeout_socket.connect.should be_false
-      end
-    end
-  end
-
-  describe "#alive?" do
-    context "when not connected" do
-      let(:socket) do
-        described_class.new("127.0.0.1", 99999).tap do |socket|
-          socket.stub(:connect)
-        end
-      end
-
-      it "should be false" do
-        socket.should_not be_alive
-      end
-    end
-
-    context "when connected but server goes away" do
-      before do
-        remote = server.accept
-        remote.shutdown
-        # Give the socket time to be notified
-        sleep 0.1
-      end
-
-      it "should be false" do
-        socket.should_not be_alive
-      end
-    end
-
-    context "when connected but server goes away" do
-      before do
-        server.close
-        # Give the socket time to be notified
-        sleep 0.1
-      end
-
-      it "should be false" do
-        socket.should_not be_alive
-      end
-    end
-
-    context "when connect is explicitly closed" do
-      before do
-        socket.close
-      end
-
-      it "should be false" do
-        socket.should_not be_alive
-      end
-    end
-
-    context "when connected and server is open" do
-      it "should be true" do
-        socket.should be_alive
-      end
-    end
-  end
-
-  describe "#execute" do
-    let(:query) { Moped::Protocol::Query.new(:moped, :test, {}) }
-
-    context "when competing threads attempt to query" do
-      let(:messages) do
-        10.times.map do |i|
-          Moped::Protocol::Insert.new(:moped, :test, {}).tap do |query|
-            query.stub(request_id: 123)
-          end
-        end
-      end
-
-      it "never issues a partial write" do
-        socket
-
-        threads = 10.times.map do |i|
-          Thread.new do
-            Thread.current.abort_on_exception = true
-            socket.execute messages[i]
-          end
-        end
-
-        threads.each(&:join)
-
-        sock = server.accept
-        data = sock.read messages.join.length
-
-        messages.each do |message|
-          fail "server received partial write" unless data.include? message.serialize
-        end
-      end
-    end
-  end
-
-  describe "#parse_reply" do
-    let(:raw) do
-      Moped::Protocol::Reply.allocate.tap do |reply|
-        reply.request_id = 1
-        reply.response_to = 1
-        reply.op_code = 1
-        reply.flags = [:await_capable]
-        reply.offset = 4
-        reply.count = 1
-        reply.documents = [{"name" => "John"}]
-      end.serialize
-    end
-
-    let(:reply) do
-      socket.parse_reply(raw.length, raw[4..-1])
-    end
-
-    it "sets the length" do
-      reply.length.should eq raw.length
-    end
-
-    it "sets the response_to" do
-      reply.response_to.should eq 1
-    end
-
-    it "sets the request id" do
-      reply.request_id.should eq 1
-    end
-
-    it "sets the flags" do
-      reply.flags.should eq [:await_capable]
-    end
-
-    it "sets the offset" do
-      reply.offset.should eq 4
-    end
-
-    it "sets the count" do
-      reply.count.should eq 1
-    end
-
-    it "sets the documents" do
-      reply.documents.should eq [{"name" => "John"}]
-    end
-  end
-
-  describe "#close" do
-    let(:exception) { RuntimeError.new }
-    let(:callback) { stub }
-
-    it "closes the connection" do
-      connection.should_receive(:close).at_least(1)
-      socket.close
-    end
-
-    it "marks the socket as dead" do
-      socket.close
-      socket.should_not be_alive
-    end
-  end
-
-  describe "#login" do
-
-    let(:connection) do
-      Support::MockConnection.new
-    end
-
-    before do
-      socket.stub(connection: connection)
-    end
-
-    context "when authentication is successful" do
-      before do
-        # getnonce
-        connection.pending_replies << Hash["nonce" => "123", "ok" => 1]
-        # authenticate
-        connection.pending_replies << Hash["ok" => 1]
-      end
-
-      it "returns true" do
-        socket.login("admin", "username", "password").should be_true
-      end
-
-      it "adds the credentials to the auth cache" do
-        socket.login(:admin, "username", "password")
-        socket.auth.should eq("admin" => ["username", "password"])
-      end
-    end
-
-    context "when a nonce fails to generate" do
-      before do
-        # getnonce
-        connection.pending_replies << Hash["ok" => 0]
-      end
-
-      it "raises an operation failure" do
-        lambda do
-          socket.login(:admin, "username", "password")
-        end.should raise_exception(Moped::Errors::OperationFailure)
-      end
-
-      it "does not add the credentials to the auth cache" do
-        socket.login(:admin, "username", "password") rescue nil
-        socket.auth.should be_empty
-      end
-    end
-
-    context "when authentication fails" do
-      before do
-        # getnonce
-        connection.pending_replies << Hash["nonce" => "123", "ok" => 1]
-        # authenticate
-        connection.pending_replies << Hash["ok" => 0]
-      end
-
-      it "raises an operation failure" do
-        lambda do
-          socket.login(:admin, "username", "password")
-        end.should raise_exception(Moped::Errors::OperationFailure)
-      end
-
-      it "does not add the credentials to the auth cache" do
-        socket.login(:admin, "username", "password") rescue nil
-        socket.auth.should be_empty
-      end
-    end
-
-  end
-
-  describe "#logout" do
-
-    let(:connection) do
-      Support::MockConnection.new
-    end
-
-    before do
-      socket.stub(connection: connection)
-      socket.auth["admin"] = ["username", "password"]
-    end
-
-    context "when logout is successful" do
-      before do
-        connection.pending_replies << Hash["ok" => 1]
-      end
-
-      it "removes the stored credentials" do
-        socket.logout :admin
-        socket.auth.should be_empty
-      end
-    end
-
-    context "when logout is unsuccessful" do
-      before do
-        connection.pending_replies << Hash["ok" => 0]
-      end
-
-      it "does not remove the stored credentials" do
-        socket.logout :admin rescue nil
-        socket.auth.should_not be_empty
-      end
-
-      it "raises an operation failure" do
-        lambda do
-          socket.logout :admin
-        end.should raise_exception(Moped::Errors::OperationFailure)
-      end
-    end
-
-  end
-
-  describe "#apply_auth" do
-    context "when the socket is unauthenticated" do
-      it "logs in with each credential provided" do
-        socket.should_receive(:login).with("admin", "username", "password")
-        socket.should_receive(:login).with("test", "username", "password")
-
-        socket.apply_auth(
-          "admin" => ["username", "password"],
-          "test" => ["username", "password"]
-        )
-      end
-    end
-
-    context "when the socket is authenticated" do
-      before do
-        socket.auth["admin"] = ["username", "password"]
-      end
-
-      context "and a credential is unchanged" do
-        it "does nothing" do
-          socket.should_not_receive(:login)
-          socket.apply_auth("admin" => ["username", "password"])
-        end
-      end
-
-      context "and a credential changes" do
-        it "logs in with the new credentials" do
-          socket.should_receive(:login).with("admin", "newuser", "password")
-          socket.apply_auth("admin" => ["newuser", "password"])
-        end
-      end
-
-      context "and a credential is removed" do
-        it "logs out from the database" do
-          socket.should_receive(:logout).with("admin")
-          socket.apply_auth({})
-        end
-      end
-
-      context "and a credential is added" do
-        it "logs in with the added credentials" do
-          socket.should_receive(:login).with("test", "username", "password")
-          socket.apply_auth(
-            "admin" => ["username", "password"],
-            "test" => ["username", "password"]
-          )
-        end
-      end
-    end
-  end
-
-  describe "instrument" do
-
-    context "when a logger is configured in debug mode" do
-      before do
-        Moped.stub(logger: mock(Logger, debug?: true))
-      end
-
-      it "logs the operations" do
-        socket.should_receive(:log_operations).once
-        socket.instrument([]) {}
-      end
-    end
-
-    context "when a logger is configured but not in debug level" do
-      before do
-        Moped.stub(logger: mock(Logger, debug?: false))
-      end
-
-      it "does not log the operations" do
-        socket.should_receive(:log_operations).never
-        socket.instrument([]) {}
-      end
-    end
-
-    context "when no logger is configured" do
-      before do
-        Moped.stub(logger: nil)
-      end
-
-      it "does not log the operations" do
-        socket.should_receive(:log_operations).never
-        socket.instrument([]) {}
-      end
-    end
-
-    context "when an error occurs" do
-      before do
-        Moped.stub(logger: mock(Logger, debug?: true))
-      end
-
-      it "does not log the operations" do
-        socket.should_receive(:log_operations).never
-
-        lambda do
-          socket.instrument([]) { raise "inner error" }
-        end.should raise_exception("inner error")
-      end
-    end
-
-  end
-
-end
diff --git a/spec/moped_spec.rb b/spec/moped_spec.rb
deleted file mode 100644
index 671fd6c..0000000
--- a/spec/moped_spec.rb
+++ /dev/null
@@ -1 +0,0 @@
-require "spec_helper"
diff --git a/spec/replset_spec.rb b/spec/replset_spec.rb
deleted file mode 100644
index 7b9c68a..0000000
--- a/spec/replset_spec.rb
+++ /dev/null
@@ -1,106 +0,0 @@
-require "spec_helper"
-
-describe "testing" do
-  let(:cluster) { Moped::Cluster.new "", false }
-  let(:socket) { Moped::Socket.new "", 99999 }
-  let(:connection) { Support::MockConnection.new }
-
-  before do
-    socket.stub(connection: connection, alive?: true)
-  end
-
-  describe "#sync_socket" do
-  end
-end
-
-__END__
-
-# sequence for single node startup:
-#
-#   connection failure (node not up)
-
-{"ismaster" => true, "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-# sequence for replica set startup pre-initialize:
-#
-#   connection failure
-
-{"ismaster" => false, "secondary" => false, "info" => "can't get local.system.replset config from self or any seed (EMPTYCONFIG)", "isreplicaset" => true, "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-# sequence for replica set startup (master):
-#
-#   connection failure (node not up)
-
-{"ismaster" => false, "secondary" => false, "info" => "can't get local.system.replset config from self or any seed (EMPTYCONFIG)", "isreplicaset" => true, "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-{"setName" => "3ff029114780", "ismaster" => false, "secondary" => true, "hosts" => ["localhost:59246", "localhost:59248", "localhost:59247"], "me" => "localhost:59246", "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-{"ismaster" => false, "secondary" => false, "info" => "Received replSetInitiate - should come online shortly.", "isreplicaset" => true, "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-{"setName" => "3ff029114780", "ismaster" => true, "secondary" => false, "hosts" => ["localhost:59246", "localhost:59248", "localhost:59247"], "primary" => "localhost:59246", "me" => "localhost:59246", "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-# sequence for replica set startup (secondary):
-#
-#   connection failure (node not up)
-
-{"ismaster" => false, "secondary" => false, "info" => "can't get local.system.replset config from self or any seed (EMPTYCONFIG)", "isreplicaset" => true, "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-{"setName" => "3fef4842b608", "ismaster" => false, "secondary" => false, "hosts" => ["localhost:61085", "localhost:61086", "localhost:61084"], "me" => "localhost:61085", "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-{"setName" => "3fef4842b608", "ismaster" => false, "secondary" => false, "hosts" => ["localhost:61085", "localhost:61086", "localhost:61084"], "primary" => "localhost:61084", "me" => "localhost:61085", "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-{"setName" => "3fef4842b608", "ismaster" => false, "secondary" => true, "hosts" => ["localhost:61085", "localhost:61086", "localhost:61084"], "primary" => "localhost:61084", "me" => "localhost:61085", "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-__END__
-
-describe Moped::Cluster do
-  context "when connecting to a single seed" do
-    context "and the seed is down"
-
-    context "and that seed is primary" do
-      it "finds the primary node"
-
-      it "adds the secondary to the dynamic seeds"
-      it "adds the arbiter to the dynamic seeds"
-    end
-
-    context "and that seed is secondary" do
-      it "finds the secondary node"
-
-      it "adds the primary to the dynamic seeds"
-      it "adds the arbiter to the dynamic seeds"
-    end
-
-    context "and that seed is an arbiter" do
-      it "adds the primary to the dynamic seeds"
-      it "adds the secondary to the dynamic seeds"
-    end
-  end
-
-  context "when connected to a single seed" do
-    context "and that seed goes down" do
-      it "is able to resync from discovered nodes"
-    end
-  end
-
-  context "when connecting to a replica set" do
-    context "and the replica set is not initiated"
-    context "and the replica set is partially initiated"
-    context "and there is no master node"
-    context "and there is no secondary node"
-  end
-
-  context "when connected to a replica set" do
-    context "and the primary node goes down" do
-      it "issues reads to the secondary"
-    end
-
-    context "and the primary node changes"
-
-    context "and the secondary node goes down" do
-      it "issues inserts to the primary"
-      it "issues reads to the primary"
-    end
-  end
-
-end
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index 24652e6..2c7a4a7 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -1,9 +1,24 @@
 require "java" if RUBY_PLATFORM == "java"
-require "bundler"
-Bundler.require
+require "rspec"
 
 $:.unshift((Pathname(__FILE__).dirname.parent + "lib").to_s)
 
 require "moped"
+require "support/mongohq"
+require "support/replica_set_simulator"
+require "support/stats"
 
-require "support/mock_connection"
+RSpec.configure do |config|
+  Support::Stats.install!
+
+  config.include Support::ReplicaSetSimulator::Helpers, replica_set: true
+
+  config.filter_run_excluding mongohq: ->(value) do
+    return true if value == :replica_set && !Support::MongoHQ.replica_set_configured?
+    return true if value == :auth && !Support::MongoHQ.auth_node_configured?
+  end
+
+  unless Support::MongoHQ.replica_set_configured? || Support::MongoHQ.auth_node_configured?
+    $stderr.puts Support::MongoHQ.message
+  end
+end
diff --git a/spec/support/mock_connection.rb b/spec/support/mock_connection.rb
deleted file mode 100644
index 9e6ac4d..0000000
--- a/spec/support/mock_connection.rb
+++ /dev/null
@@ -1,39 +0,0 @@
-module Support
-  class MockConnection
-
-    attr_reader :pending_replies
-
-    def initialize
-      @connected = true
-      @buffer = StringIO.new
-      @pending_replies = []
-    end
-
-    def write(*args)
-
-    end
-
-    def read(*args)
-      while documents = pending_replies.shift
-        reply = Moped::Protocol::Reply.allocate.tap do |reply|
-          documents = [documents] unless documents.is_a? Array
-          reply.documents = documents
-          reply.count = documents.length
-        end
-
-        reply.serialize(@buffer.string)
-      end
-
-      @buffer.read(*args)
-    end
-
-    def closed?
-      !@connected
-    end
-
-    def close
-      @connected = false
-    end
-
-  end
-end
diff --git a/spec/support/mongohq.rb b/spec/support/mongohq.rb
new file mode 100644
index 0000000..706fdcb
--- /dev/null
+++ b/spec/support/mongohq.rb
@@ -0,0 +1,61 @@
+module Support
+  module MongoHQ
+    extend self
+
+    def replica_set_configured?
+      ENV["MONGOHQ_REPL_PASS"]
+    end
+
+    def replica_set_seeds
+      [ENV["MONGOHQ_REPL_1_URL"], ENV["MONGOHQ_REPL_2_URL"]]
+    end
+
+    def replica_set_credentials
+      [ENV["MONGOHQ_REPL_USER"], ENV["MONGOHQ_REPL_PASS"]]
+    end
+
+    def replica_set_database
+      ENV["MONGOHQ_REPL_NAME"]
+    end
+
+    def replica_set_session(auth = true)
+      session = Moped::Session.new replica_set_seeds, database: replica_set_database
+      session.login *replica_set_credentials if auth
+      session
+    end
+
+    def auth_seeds
+      [ENV["MONGOHQ_SINGLE_URL"]]
+    end
+
+    def auth_node_configured?
+      ENV["MONGOHQ_SINGLE_PASS"]
+    end
+
+    def auth_credentials
+      [ENV["MONGOHQ_SINGLE_USER"], ENV["MONGOHQ_SINGLE_PASS"]]
+    end
+
+    def auth_database
+      ENV["MONGOHQ_SINGLE_NAME"]
+    end
+
+    def auth_session(auth = true)
+      session = Moped::Session.new auth_seeds, database: auth_database
+      session.login *auth_credentials if auth
+      session
+    end
+
+    def message
+      %Q{
+      ---------------------------------------------------------------------
+      Moped runs specs for authentication and replica sets against MongoHQ.
+
+      If you want to run these specs and need the credentials, contact
+      durran at gmail dot com.
+      ---------------------------------------------------------------------
+      }
+    end
+
+  end
+end
diff --git a/spec/support/replica_set_simulator.rb b/spec/support/replica_set_simulator.rb
new file mode 100644
index 0000000..eb09a60
--- /dev/null
+++ b/spec/support/replica_set_simulator.rb
@@ -0,0 +1,306 @@
+module Support
+
+  # This is a helper class for testing replica sets. It works by starting up a
+  # TCP server socket for each desired node. It then proxies all traffic
+  # between a real mongo instance and the client app, with the exception of
+  # ismaster commands, which it returns simulated responses for.
+  class ReplicaSetSimulator
+
+    module Helpers
+      def self.included(context)
+        context.before :all do
+          @replica_set = ReplicaSetSimulator.new
+          @replica_set.start
+
+          @primary, @secondaries = @replica_set.initiate
+        end
+
+        context.after :all do
+          @replica_set.stop
+        end
+
+        context.after :each do
+          @replica_set.nodes.each &:restart
+        end
+
+        context.let :seeds do
+          @replica_set.nodes.map &:address
+        end
+      end
+    end
+
+    attr_reader :nodes
+    attr_reader :manager
+
+    def initialize(nodes = 3)
+      @nodes = nodes.times.map { Node.new(self) }
+      @manager = ConnectionManager.new(@nodes)
+      @mongo = TCPSocket.new "127.0.0.1", 27017
+    end
+
+    # Start the mock replica set.
+    def start
+      @nodes.each &:start
+      @worker = Thread.start do
+        Thread.abort_on_exception = true
+        catch(:shutdown) do
+          loop do
+            Moped.logger.debug "replica_set: waiting for next client"
+            server, client = @manager.next_client
+
+            if server
+              Moped.logger.debug "replica_set: proxying incoming request to mongo"
+              server.proxy(client, @mongo)
+            else
+              Moped.logger.debug "replica_set: no requests; passing"
+              Thread.pass
+            end
+          end
+        end
+      end
+    end
+
+    # Pick a node to be master, and mark the rest as secondary
+    def initiate
+      primary, *secondaries = @nodes.shuffle
+
+      primary.promote
+      secondaries.each &:demote
+
+      return primary, secondaries
+    end
+
+    # Shut down the mock replica set.
+    def stop
+      @manager.shutdown
+      @nodes.each &:stop
+    end
+
+    class Node
+
+      attr_reader :port
+      attr_reader :host
+
+      def initialize(set)
+        @set = set
+        @primary = false
+        @secondary = false
+
+        server = TCPServer.new 0
+        @host = Socket.gethostname
+        @port = server.addr[1]
+        server.close
+      end
+
+      def ==(other)
+        @host == other.host && @port == other.port
+      end
+
+      def address
+        "#{@host}:#{@port}"
+      end
+
+      def primary?
+        @primary
+      end
+
+      def secondary?
+        @secondary
+      end
+
+      def status
+        {
+          "ismaster" => @primary,
+          "secondary" => @secondary,
+          "hosts" => @set.nodes.map(&:address),
+          "me" => address,
+          "maxBsonObjectSize" => 16777216,
+          "ok" => 1.0
+        }
+      end
+
+      def status_reply
+        reply = Moped::Protocol::Reply.new
+        reply.count = 1
+        reply.documents = [status]
+        reply
+      end
+
+      OP_QUERY = 2004
+      OP_GETMORE = 2005
+
+      # Stop and start the node.
+      def restart
+        stop
+        start
+      end
+
+      # Start the node.
+      def start
+        @server = TCPServer.new @port
+      end
+
+      # Stop the node.
+      def stop
+        if @server
+          hiccup
+
+          @server.close
+          @server = nil
+        end
+      end
+      alias close stop
+
+      def accept
+        to_io.accept
+      end
+
+      def closed?
+        !@server || @server.closed?
+      end
+
+      def to_io
+        @server
+      end
+
+      # Mark this node as secondary.
+      def demote
+        @primary = false
+        @secondary = true
+
+        hiccup
+      end
+
+      def hiccup
+        @set.manager.close_clients_for(self)
+      end
+
+      # Mark this node as primary. This also closes any open connections.
+      def promote
+        @primary = true
+        @secondary = false
+
+        hiccup
+      end
+
+      # Proxies a single message from client to the mongo connection.
+      def proxy(client, mongo)
+        incoming_message = client.read(16)
+        length, op_code = incoming_message.unpack("l<x8l<")
+        incoming_message << client.read(length - 16)
+
+        if op_code == OP_QUERY && ismaster_command?(incoming_message)
+          # Intercept the ismaster command and send our own reply.
+          client.write status_reply
+        else
+          # This is a normal command, so proxy it to the real mongo instance.
+          mongo.write incoming_message
+
+          if op_code == OP_QUERY || op_code == OP_GETMORE
+            outgoing_message = mongo.read(4)
+            length, = outgoing_message.unpack('l<')
+            outgoing_message << mongo.read(length - 4)
+
+            client.write outgoing_message
+          end
+        end
+      end
+
+      private
+
+      # Checks a message to see if it's an `ismaster` query.
+      def ismaster_command?(incoming_message)
+        data = StringIO.new(incoming_message)
+        data.read(20) # header and flags
+        data.gets("\x00") # collection name
+        data.read(8) # skip/limit
+
+        selector = Moped::BSON::Document.deserialize(data)
+        selector == { "ismaster" => 1 }
+      end
+    end
+
+    class ConnectionManager
+
+      def initialize(servers)
+        @timeout = 0.1
+        @servers = servers
+        @clients = []
+      end
+
+      def shutdown
+        @servers.each &:close
+        @clients.each &:close
+        @shutdown = true
+      end
+
+      def next_client
+        throw :shutdown if @shutdown
+
+        begin
+          servers = @servers.reject &:closed?
+          clients =  @clients.reject &:closed?
+          Moped.logger.debug "replica_set: selecting on connections"
+          readable, _, errors = Kernel.select(servers + clients, nil, clients, @timeout)
+        rescue IOError, Errno::EBADF, TypeError
+          # Looks like we hit a bad file descriptor or closed connection.
+          Moped.logger.debug "replica_set: io error, retrying"
+          retry
+        end
+
+        return unless readable || errors
+
+        errors.each do |client|
+          client.close
+          @clients.delete client
+        end
+
+        clients, servers = readable.partition { |s| s.class == TCPSocket }
+
+        servers.each do |server|
+          Moped.logger.debug "replica_set: accepting new client for #{server.port}"
+          @clients << server.accept
+        end
+
+        Moped.logger.debug "replica_set: closing dead clients"
+        closed, open = clients.partition &:eof?
+        closed.each { |client| @clients.delete client }
+
+        if client = open.shift
+          Moped.logger.debug "replica_set: finding server for client"
+          server = lookup_server(client)
+
+          Moped.logger.debug "replica_set: sending client #{client.inspect} to #{server.port}"
+          return server, client
+        else
+          nil
+        end
+      end
+
+      def close_clients_for(server)
+        Moped.logger.debug "replica_set: closing open clients on #{server.port}"
+
+        @clients.reject! do |client|
+          port = client.addr(false)[1]
+
+          if port == server.port
+            client.close
+            true
+          else
+            false
+          end
+        end
+      end
+
+      def lookup_server(client)
+        port = client.addr(false)[1]
+
+        @servers.find do |server|
+          server.to_io && server.to_io.addr[1] == port
+        end
+      end
+
+    end
+
+  end
+end
diff --git a/spec/support/stats.rb b/spec/support/stats.rb
new file mode 100644
index 0000000..af8cb79
--- /dev/null
+++ b/spec/support/stats.rb
@@ -0,0 +1,51 @@
+module Support
+
+  # Module for recording operations.
+  #
+  #   Support::Stats.install!
+  #
+  #   stats = Support::Stats.collect do
+  #     session.with(safe: true)[:users].insert({})
+  #   end
+  #
+  #   ops = stats["127.0.0.1:27017"]
+  #   ops.size # => 2
+  #   ops[0].class # => Moped::Protocol::Insert
+  #   ops[1].class # => Moped::Protocol::Command
+  #
+  module Stats
+    extend self
+
+    def record(node, operations)
+      key = if node.primary?
+        :primary
+      elsif node.secondary?
+        :secondary
+      else
+        :other
+      end
+
+      @stats[key].concat(operations) if @stats
+    end
+
+    def collect
+      @stats = Hash.new { |hash, key| hash[key] = [] }
+      yield
+      @stats
+    ensure
+      @stats = nil
+    end
+
+    def install!
+      Moped::Node.class_eval <<-EOS
+        alias _logging logging
+
+        def logging(operations, &block)
+          Support::Stats.record(self, operations)
+          _logging(operations, &block)
+        end
+      EOS
+    end
+
+  end
+end

Vulnerable code: require "moped/server"
require "moped/socket"
      attr_reader :data

          data = []
          new data
          !!str.match(/^[0-9a-f]{24}$/i)
      end
      def initialize(data = nil, time = nil)
        if data
          @data = data
        elsif time
          @data = @@generator.generate(time.to_i)
        else
          @data = @@generator.next
        @@string_format % data
        Time.at(@data.pack("C4").unpack("N")[0]).utc
          new io.read(12).unpack('C*')

        io << data.pack('C12')
          @machine_id = Digest::MD5.digest(Socket.gethostname).unpack("C3")
          @last_timestamp = nil
        # Return object id data based on the current time, incrementing a
        # counter for object ids generated in the same second.
          now = Time.new.to_i

          counter = @mutex.synchronize do
            last_timestamp, @last_timestamp = @last_timestamp, now

            if last_timestamp == now
              @counter += 1
            else
              @counter = 0
            end
          generate(now, counter)
        # Generate object id data for a given time using the provided +inc+.
        def generate(time, inc = 0)
          pid = Process.pid % 0xFFFF

          [
            time >> 24 & 0xFF, # 4 bytes time (network order)
            time >> 16 & 0xFF,
            time >> 8  & 0xFF,
            time       & 0xFF,
            @machine_id[0],   # 3 bytes machine
            @machine_id[1],
            @machine_id[2],
            pid  >> 8  & 0xFF, # 2 bytes process id
            pid        & 0xFF,
            inc  >> 16 & 0xFF, # 3 bytes increment
            inc  >> 8  & 0xFF,
            inc        & 0xFF,
          ]
  # @api private
  #
  # The internal class managing connections to both a single node and replica
  # sets.
  #
  # @note Though the socket class itself *is* threadsafe, the cluster presently
  #   is not. This means that in the course of normal operations sessions can be
  #   shared across threads, but in failure modes (when a resync is required),
  #   things can possibly go wrong.
    # @return [Array] the user supplied seeds
    # @return [Boolean] whether this is a direct connection
    attr_reader :direct

    # @return [Array] all available nodes
    attr_reader :servers

    # @return [Array] seeds gathered from cluster discovery
    attr_reader :dynamic_seeds

    # @param [Array] seeds an array of host:port pairs
    # @param [Boolean] direct (false) whether to connect directly to the hosts
    # provided or to find additional available nodes.
    def initialize(seeds, direct = false)
      @seeds  = seeds
      @direct = direct

      @servers = []
      @dynamic_seeds = []
    end

    # @return [Array] available secondary nodes
    def secondaries
      servers.select(&:secondary?)
    end

    # @return [Array] available primary nodes
    def primaries
      servers.select(&:primary?)
    end

    # @return [Array] all known addresses from user supplied seeds, dynamically
    # discovered seeds, and active servers.
    def known_addresses
      [].tap do |addresses|
        addresses.concat seeds
        addresses.concat dynamic_seeds
        addresses.concat servers.map { |server| server.address }
      end.uniq
    end

    def remove(server)
      servers.delete(server)
    def reconnect
      @servers = servers.map { |server| Server.new(server.address) }
    end

    def sync
      known = known_addresses.shuffle
      seen  = {}

      sync_seed = ->(seed) do
        server = Server.new seed

        unless seen[server.resolved_address]
          seen[server.resolved_address] = true

          hosts = sync_server(server)

          hosts.each do |host|
            sync_seed[host]
          end
      known.each do |seed|
        sync_seed[seed]
      end

      unless servers.empty?
        @dynamic_seeds = servers.map(&:address)
      end

      true
    def sync_server(server)
      [].tap do |hosts|
        socket = server.socket

        if socket.connect
          info = socket.simple_query Protocol::Command.new(:admin, ismaster: 1)

          if info["ismaster"]
            server.primary = true
          end
          if info["secondary"]
            server.secondary = true
          end
          if info["primary"]
            hosts.push info["primary"]
          end

          if info["hosts"]
            hosts.concat info["hosts"]
          end
          if info["passives"]
            hosts.concat info["passives"]

          merge(server)

      end.uniq
    end

    def merge(server)
      previous = servers.find { |other| other == server }
      primary = server.primary?
      secondary = server.secondary?
      if previous
        previous.merge(server)
        servers << server
    # @param [:read, :write] mode the type of socket to return
    # @return [Socket] a socket valid for +mode+ operations
    def socket_for(mode)
      sync unless primaries.any? || (secondaries.any? && mode == :read)

      server = nil
      while primaries.any? || (secondaries.any? && mode == :read)
        if mode == :write || secondaries.empty?
          server = primaries.sample
        else
          server = secondaries.sample
        end

        if server
          socket = server.socket
          socket.connect unless socket.connection

          if socket.alive?
            break server
          else
            remove server
          end
      unless server
        raise Errors::ConnectionFailure.new("Could not connect to any primary or secondary servers")

      socket = server.socket
      socket.apply_auth auth
      socket
    # Log in to +database+ with +username+ and +password+. Does not perform the
    # actual log in, but saves the credentials for later authentication on a
    # socket.
    def login(database, username, password)
      auth[database.to_s] = [username, password]
    end
    # Log out of +database+. Does not perform the actual log out, but will log
    # out when the socket is used next.
    def logout(database)
      auth.delete(database.to_s)

      insert = Protocol::Insert.new(database.name, name, documents)
        session.execute insert

      @query_op = query_operation.dup
      @get_more_op = Protocol::GetMore.new(
        @query_op.database,
        @query_op.collection,
        0,
        @query_op.limit
      )
      @kill_cursor_op = Protocol::KillCursors.new([0])
      documents = query @query_op
        return kill if limited? && @get_more_op.limit <= 0
        documents = query @get_more_op
    def query(operation)
      reply = session.query operation
      @get_more_op.limit -= reply.count if limited?
      @get_more_op.cursor_id = reply.cursor_id
      @kill_cursor_op.cursor_ids = [reply.cursor_id]
      @query_op.limit > 0
      @get_more_op.cursor_id != 0
      session.execute kill_cursor_op
      command dropDatabase: 1
      session.cluster.login(name, username, password)
      session.cluster.logout(name)
      operation = Protocol::Command.new(name, command)

      result = session.with(consistency: :strong) do |session|
        session.simple_query(operation)
      end

      raise Errors::OperationFailure.new(
        operation, result
      ) unless result["ok"] == 1.0

      result

  # The namespace for all errors generated by Moped.
    # Raised when providing an invalid string from an object id.
    class InvalidObjectId < StandardError
      def initialize(string)
        super("'#{string}' is not a valid object id.")
      end
    end
    class MongoError < StandardError; end

    # Exception class for exceptions generated as a direct result of an
    # operation, such as a failed insert or an invalid command.
    class OperationFailure < MongoError

    # A special kind of OperationFailure, raised when Mongo sets the
    # :query_failure flag on a query response.
    class QueryFailure < OperationFailure; end
      def initialize(database, command)
        super database, :$cmd, command, limit: -1
      session.simple_query(operation)
      cursor = Cursor.new(session.with(retain_socket: true), operation)
      update = Protocol::Update.new(
        operation.database,
        operation.collection,
        operation.selector,
        change,
        flags: flags
      )

        session.execute update
      delete = Protocol::Delete.new(
        operation.database,
        operation.collection,
        operation.selector,
        flags: [:remove_first]
      )

        session.execute delete
      delete = Protocol::Delete.new(
        operation.database,
        operation.collection,
        operation.selector
      )

        session.execute delete
module Moped

  # @api private
  #
  # The internal class for storing information about a server.
  class Server

    # @return [String] the original host:port address provided
    attr_reader :address

    # @return [String] the resolved host:port address
    attr_reader :resolved_address

    # @return [String] the resolved ip address
    attr_reader :ip_address

    # @return [Integer] the resolved port
    attr_reader :port

    attr_writer :primary
    attr_writer :secondary

    def initialize(address)
      @address = address

      host, port = address.split(":")
      port = port ? port.to_i : 27017

      ip_address = ::Socket.getaddrinfo(host, nil, ::Socket::AF_INET, ::Socket::SOCK_STREAM).first[3]

      @primary = @secondary = false
      @ip_address = ip_address
      @port = port
      @resolved_address = "#{ip_address}:#{port}"
    end

    def primary?
      !!@primary
    end

    def secondary?
      !!@secondary
    end

    def merge(other)
      @primary = other.primary?
      @secondary = other.secondary?

      other.close
    end

    def close
      if @socket
        @socket.close
        @socket = nil
      end
    end

    def socket
      @socket ||= Socket.new(ip_address, port)
    end

    def ==(other)
      self.class === other && hash == other.hash
    end
    alias eql? ==

    def hash
      [ip_address, port].hash
    end

  end
end
      @cluster = Cluster.new(seeds)
    #   session[:people].     john, mary = session[:people].find.one # => { :name => "John" }
      session.cluster.reconnect
    # @private
    def current_database
      return @current_database if defined? @current_database

      if database = options[:database]
        set_current_database(database)
      else
        raise "No database set for session. Call #use or #with before accessing the database"
      end
    end

    # @private
    def simple_query(query)
      query.limit = -1

      query(query).documents.first
    end

    # @private
    def query(query)
      if options[:consistency] == :eventual
        query.flags |= [:slave_ok] if query.respond_to? :flags
        mode = :read
      else
        mode = :write
      end

      reply = socket_for(mode).execute(query)

      reply.tap do |reply|
        if reply.flags.include?(:query_failure)
          raise Errors::QueryFailure.new(query, reply.documents.first)
        end
      end
    end

    # @private
    def execute(op)
      mode = options[:consistency] == :eventual ? :read : :write
      socket = socket_for(mode)

      if safe?
        last_error = Protocol::Command.new(
          "admin", { getlasterror: 1 }.merge(safety)
        )

        socket.execute(op, last_error).documents.first.tap do |result|
          raise Errors::OperationFailure.new(
            op, result
          ) if result["err"] || result["errmsg"]
        end
      else
        socket.execute(op)
      end
    end

    private

    def socket_for(mode)
      if options[:retain_socket]
        @socket ||= cluster.socket_for(mode)
        cluster.socket_for(mode)
    def dup
      session = super
      session.instance_variable_set :@options, options.dup
        session.send(:remove_instance_variable, :@current_database)

      session
require "timeout"

module Moped

  # @api private
  #
  # The internal class wrapping a socket connection.
  class Socket

    # Thread-safe atomic integer.
    class RequestId
      def initialize
        @mutex = Mutex.new
        @id = 0
      end

      def next
        @mutex.synchronize { @id += 1 }
      end
    end

    attr_reader :connection

    attr_reader :host
    attr_reader :port

    def initialize(host, port)
      @host = host
      @port = port

      @mutex = Mutex.new
      @request_id = RequestId.new
    end

    # @return [true, false] whether the connection was successful
    # @note The connection timeout is currently just 0.5 seconds, which should
    #   be sufficient, but may need to be raised or made configurable for
    #   high-latency situations. That said, if connecting to the remote server
    #   takes that long, we may not want to use the node any way.
    def connect
      return true if connection

      Timeout::timeout 0.5 do
        @connection = TCPSocket.new(host, port)
      end
    rescue Errno::ECONNREFUSED, Timeout::Error
      return false
    end

    # @return [true, false] whether this socket connection is alive
    def alive?
      if connection
        return false if connection.closed?

        @mutex.synchronize do
          if select([connection], nil, nil, 0)
            !connection.eof? rescue false
          else
            true
          end
        end
      else
        false
      end
    end

    # Execute the operations on the connection.
    def execute(*ops)
      instrument(ops) do
        buf = ""

        last = ops.each do |op|
          op.request_id = @request_id.next
          op.serialize buf
        end.last

        if Protocol::Query === last || Protocol::GetMore === last
          length = nil

          @mutex.synchronize do
            connection.write buf

            length, = connection.read(4).unpack('l<')

            # Re-use the already allocated buffer used for writing the command.
            connection.read(length - 4, buf)
          end

          parse_reply length, buf
        else
          @mutex.synchronize do
            connection.write buf
          end

          nil
        end
      end
    end

    def parse_reply(length, data)
      buffer = StringIO.new data

      reply = Protocol::Reply.allocate

      reply.length = length

      reply.request_id,
        reply.response_to,
        reply.op_code,
        reply.flags,
        reply.cursor_id,
        reply.offset,
        reply.count = buffer.read(32).unpack('l4<q<l2<')

      reply.documents = reply.count.times.map do
        BSON::Document.deserialize(buffer)
      end

      reply
    end

    # Executes a simple (one result) query and returns the first document.
    #
    # @return [Hash] the first document in a result set.
    def simple_query(query)
      query = query.dup
      query.limit = -1

      execute(query).documents.first
    end

    # Manually closes the connection
    def close
      @mutex.synchronize do
        connection.close if connection && !connection.closed?
        @connection = nil
      end
    end

    def auth
      @auth ||= {}
    end

    def apply_auth(credentials)
      return if auth == credentials
      logouts = auth.keys - credentials.keys

      logouts.each do |database|
        logout database
      end

      credentials.each do |database, (username, password)|
        login(database, username, password) unless auth[database] == [username, password]
      end
    end

    def login(database, username, password)
      getnonce = Protocol::Command.new(database, getnonce: 1)
      result = simple_query getnonce

      raise Errors::OperationFailure.new(getnonce, result) unless result["ok"] == 1

      authenticate = Protocol::Commands::Authenticate.new(database, username, password, result["nonce"])
      result = simple_query authenticate
      raise Errors::OperationFailure.new(authenticate, result) unless result["ok"] == 1

      auth[database.to_s] = [username, password]
    end

    def logout(database)
      command = Protocol::Command.new(database, logout: 1)
      result = simple_query command
      raise Errors::OperationFailure.new(command, result) unless result["ok"] == 1
      auth.delete(database.to_s)
    end

    def instrument(ops)
      instrument_start = (logger = Moped.logger) && logger.debug? && Time.now
      yield
    ensure
      log_operations(logger, ops, Time.now - instrument_start) if instrument_start && !$!
    end

    def log_operations(logger, ops, duration)
      prefix  = "  MOPED: #{host}:#{port} "
      indent  = " "*prefix.length
      runtime = (" (%.1fms)" % duration)

      if ops.length == 1
        logger.debug prefix + ops.first.log_inspect + runtime
      else
        first, *middle, last = ops

        logger.debug prefix + first.log_inspect
        middle.each { |m| logger.debug indent + m.log_inspect }
        logger.debug indent + last.log_inspect + runtime
      end
    end

  end
end
      1000.times do
        session[:people].insert(name: "John")
        session[:people].find.one
      1000.times do
        session[:people].insert(name: "John")
        session[:people].find.one
require "spec_helper"

describe Moped::Protocol do
  Protocol = Moped::Protocol

  let(:connection) do
    TCPSocket.new("localhost", 27017)
  end

  after do
    connection.close unless connection.closed?
  end

  describe "authentication" do
    context "when nonce is invalid" do
      let(:auth) do
        Protocol::Commands::Authenticate.new :admin, "user", "pass", "fakenonce"
      end

      it "fails" do
        connection.write auth
        reply = Protocol::Reply.deserialize(connection).documents[0]
        reply["ok"].should eq 0.0
      end
    end

    context "when nonce is valid but user doesn't exist" do
      let(:nonce) do
        command = Protocol::Command.new :admin, getnonce: 1
        connection.write command
        Protocol::Reply.deserialize(connection).documents[0]["nonce"]
      end

      let(:auth) do
        Protocol::Commands::Authenticate.new :admin, "user", "pass", nonce
      end

      it "fails" do
        connection.write auth
        reply = Protocol::Reply.deserialize(connection).documents[0]
        reply["ok"].should eq 0.0
      end
    end

    context "when nonce is valid but password is wrong" do
      let(:nonce) do
        command = Protocol::Command.new "moped-protocol-spec", getnonce: 1
        connection.write command
        Protocol::Reply.deserialize(connection).documents[0]["nonce"]
      end

      let(:auth) do
        Protocol::Commands::Authenticate.new "moped-protocol-spec",
          "moped",
          "pass",
          nonce
      end

      before do
        connection.write Protocol::Insert.new(
          "moped-protocol-spec",
          "system.users",
          [{ user: "moped", pwd: Digest::MD5.hexdigest("moped:mongo:password") }]
        )
      end

      it "fails" do
        connection.write auth
        reply = Protocol::Reply.deserialize(connection).documents[0]
        reply["ok"].should eq 0.0
      end
    end

    context "when authentication is valid" do
      let(:nonce) do
        command = Protocol::Command.new "moped-protocol-spec", getnonce: 1
        connection.write command
        Protocol::Reply.deserialize(connection).documents[0]["nonce"]
      end

      let(:auth) do
        Protocol::Commands::Authenticate.new "moped-protocol-spec",
          "moped",
          "password",
          nonce
      end

      before do
        connection.write Protocol::Insert.new(
          "moped-protocol-spec",
          "system.users",
          [{ user: "moped", pwd: Digest::MD5.hexdigest("moped:mongo:password") }]
        )
      end

      it "succeeds" do
        connection.write auth
        reply = Protocol::Reply.deserialize(connection).documents[0]
        reply["ok"].should eq 1.0
      end
    end
  end
end
require "spec_helper"

describe Moped::Protocol do
  let(:Protocol) { Moped::Protocol }

  let(:connection) do
    TCPSocket.new("localhost", 27017)
  end

  after do
    connection.close unless connection.closed?
  end

  describe "reply response flags" do
    let(:reply) do
      Protocol::Reply.deserialize(connection)
    end

    context "when get more is called with an invalid cursor" do
      let(:get_more) do
        Protocol::GetMore.new("moped-protocol", "suite", 0, 0)
      end

      specify "the cursor not found flag is set" do
        connection.write get_more
        reply.flags.should include :cursor_not_found
      end
    end

    context "when query generates an error" do
      let(:query) do
        Protocol::Query.new "moped-protocol", "people", { '$in' => 1 }, limit: -1
      end

      specify "the query failure flag is set" do
        connection.write query
        reply.flags.should eq [:query_failure]
      end
    end

    context "when mongod supports await data query option" do
      let(:query) do
        Protocol::Query.new "admin", "$cmd", { buildinfo: 1 }, limit: -1
      end

      specify "the await capable flag is set" do
        connection.write query

        if reply.documents[0]["version"] >= "1.6"
          reply.flags.should eq [:await_capable]
        else
          reply.flags.should eq []
        end
      end
    end

  end
end
# encoding: utf-8

require "spec_helper"

describe Moped::Session do
  context "with a single master node" do
    let(:session) { Moped::Session.new ["127.0.0.1:27017"], database: "moped_test" }

    after do
      session[:people].drop if session[:people].find.count > 0
      session.cluster.servers.each(&:close)
    end

    it "inserts and queries a single document" do
      id = Moped::BSON::ObjectId.new
      session[:people].insert(_id: id, name: "John")
      john = session[:people].find(_id: id).one
      john["_id"].should eq id
      john["name"].should eq "John"
    end

    it "inserts and queries on utf-8 data" do
      id = Moped::BSON::ObjectId.new
      doc = {
        "_id" => id,
        "gültig" => "1",
        "1" => "gültig",
        "2" => :"gültig",
        "3" => ["gültig"],
        "4" => /gültig/
      }
      session[:people].insert(doc)
      session[:people].find(_id: id).one.should eq doc
    end

    it "can explain a query" do
      id = Moped::BSON::ObjectId.new
      session[:people].find(_id: id).explain["cursor"].should eq("BasicCursor")
    end

    it "can explain a query with a sort" do
      id = Moped::BSON::ObjectId.new
      query = session[:people].find(_id: id)
      query.sort(_id: 1).explain["cursor"].should eq("BasicCursor")
    end

    it "drops a collection" do
      session.command(count: :people)["n"].should eq 0
      session[:people].insert(name: "John")
      session.command(count: :people)["n"].should eq 1
      session[:people].drop
      session.command(count: :people)["n"].should eq 0
    end

    it "can be inserted into safely" do
      session.with(safe: true) do |session|
        session[:people].insert(name: "John")["ok"].should eq 1
      end
    end

    it "raises an error on a failed insert in safe mode" do
      session.with(safe: true) do |session|
        lambda do
          session[:people].insert("$invalid" => nil)
        end.should raise_exception(Moped::Errors::OperationFailure)
      end
    end

    it "can sort documents" do
      session[:people].insert([{name: "John"}, {name: "Mary"}])
      session[:people].find.sort(_id: -1).first["name"].should eq "Mary"
    end

    it "can update documents" do
      id = Moped::BSON::ObjectId.new
      session[:people].insert(_id: id, name: "John")
      mary = session[:people].find(_id: id).one
      mary["name"].should eq "John"
      session[:people].find(_id: id).update(name: "Mary")
      mary = session[:people].find(_id: id).one
      mary["_id"].should eq id
      mary["name"].should eq "Mary"
    end

    it "can update documents safely" do
      id = Moped::BSON::ObjectId.new
      session[:people].insert(_id: id, name: "John")
      mary = session[:people].find(_id: id).one
      mary["name"].should eq "John"
      session.with(safe: true) do |session|
        session[:people].find(_id: id).update(name: "Mary")["ok"].should eq 1
      end
    end

    it "can update multiple documents" do
      session[:people].insert([{name: "John"}, {name: "Mary"}])
      session[:people].find.update_all("$set" => { "last_name" => "Unknown" })

      session[:people].find.sort(_id: -1).first["last_name"].should eq "Unknown"
      session[:people].find.sort(_id: 1).first["last_name"].should eq "Unknown"
    end

    it "can upsert documents" do
      session[:people].find.upsert(name: "Mary")
      mary = session[:people].find(name: "Mary").one
      mary["name"].should eq "Mary"
    end

    it "can delete a single document" do
      session[:people].insert([{name: "John"}, {name: "John"}])
      session[:people].find(name: "John").remove
      session[:people].find.count.should eq 1
    end

    it "can delete a single document safely" do
      session[:people].insert([{name: "John"}, {name: "John"}])
      session.with(safe: true) do |session|
        session[:people].find(name: "John").remove["ok"].should eq 1
      end
      session[:people].find.count.should eq 1
    end

    it "can delete a multiple documents" do
      session[:people].insert([{name: "John"}, {name: "John"}])
      session[:people].find(name: "John").remove_all
      session[:people].find.count.should eq 0
    end

    it "can retrieve multiple documents with fixed limit" do
      session[:people].insert([{name: "John"}, {name: "Mary"}])
      john, mary = session[:people].find.limit(-2).sort(name: 1).to_a
      john["name"].should eq "John"
      mary["name"].should eq "Mary"
    end

    it "can retrieve distinct values" do
      session[:people].insert([{name: "John"}, {name: "Mary"}])
      values = session[:people].find.distinct(:name)
      values.should eq [ "John", "Mary" ]
    end

    it "can retrieve no documents" do
      session[:people].find.limit(-2).sort(name: 1).to_a.should eq []
    end

    it "can limit a result set" do
      documents = 100.times.map { { _id: Moped::BSON::ObjectId.new } }
      session[:people].insert(documents)
      session[:people].find.limit(20).to_a.length.should eq 20
    end

    it "does not leave open cursors" do
      documents = 100.times.map { { _id: Moped::BSON::ObjectId.new } }
      session[:people].insert(documents)
      session[:people].find.limit(20).to_a.length.should eq 20
      status = session.command serverStatus: 1
      status["cursors"]["totalOpen"].should eq 0
    end

    it "can retrieve large result sets" do
      documents = 1000.times.map do
        { _id: Moped::BSON::ObjectId.new }
      end
      session[:people].insert(documents)
      session[:people].find.to_a.length.should eq 1000
    end

    it "can have multiple connections" do
      status = session.command serverStatus: 1
      count = status["connections"]["current"]
      new_session = session.new
      status = new_session.command serverStatus: 1
      status["connections"]["current"].should eq count + 1
    end
  end
end
    [78, 77, 102, 52, 59, 57, 182, 132, 7, 0, 0, 1]
      it "initializes with the strings bytes" do
        Moped::BSON::ObjectId.should_receive(:new).with(bytes)
  describe "#initialize" do

    context "with data" do
      it "sets the object id's data" do
        Moped::BSON::ObjectId.new(bytes).data.should == bytes
      end

    context "with a time" do
      it "sets the generation time" do
        time = Time.at((Time.now.utc - 64800).to_i).utc
        Moped::BSON::ObjectId.new(nil, time).generation_time.should == time
      end
    end

        Moped::BSON::ObjectId.new(bytes).should == Moped::BSON::ObjectId.new(bytes)
        Moped::BSON::ObjectId.new(bytes).should eql Moped::BSON::ObjectId.new(bytes)
        Moped::BSON::ObjectId.new(bytes).hash.should eq Moped::BSON::ObjectId.new(bytes).hash
      Moped::BSON::ObjectId.new(bytes).to_s.should eq "4e4d66343b39b68407000001"
describe Moped::Cluster do

  let(:master) do
    TCPServer.new "127.0.0.1", 0
  end

  let(:secondary_1) do
    TCPServer.new "127.0.0.1", 0
  end

  let(:secondary_2) do
    TCPServer.new "127.0.0.1", 0
  end

  describe "initialize" do
    let(:cluster) do
      Moped::Cluster.new(["127.0.0.1:27017","127.0.0.1:27018"], true)
    end

    it "stores the list of seeds" do
      cluster.seeds.should eq ["127.0.0.1:27017", "127.0.0.1:27018"]
    end

    it "stores whether the connection is direct" do
      cluster.direct.should be_true
    end

    it "has an empty list of primaries" do
      cluster.primaries.should be_empty
    end

    it "has an empty list of secondaries" do
      cluster.secondaries.should be_empty
    end

    it "has an empty list of servers" do
      cluster.servers.should be_empty
    end

    it "has an empty list of dynamic seeds" do
      cluster.dynamic_seeds.should be_empty
    end
  describe "#sync" do
    let(:cluster) { Moped::Cluster.new(["127.0.0.1:27017"]) }

    it "syncs each seed node" do
      server = Moped::Server.allocate
      Moped::Server.should_receive(:new).with("127.0.0.1:27017").and_return(server)

      cluster.should_receive(:sync_server).with(server).and_return([])
      cluster.sync
    end
  end

  describe "#sync_server" do
    let(:cluster) { Moped::Cluster.new [""], false }
    let(:server) { Moped::Server.new("localhost:27017") }
    let(:socket) { Moped::Socket.new "", 99999 }
    let(:connection) { Support::MockConnection.new }

    before do
      socket.stub(connection: connection, alive?: true)
      server.stub(socket: socket)
    context "when node is not running" do
      it "returns nothing" do
        socket.stub(connect: false)

        cluster.sync_server(server).should be_empty
    context "when talking to a single node" do
        connection.pending_replies << Hash[
          "ismaster" => true,
          "maxBsonObjectSize" => 16777216,
          "ok" => 1.0
        ]
      end

      it "adds the node to the master set" do
        cluster.sync_server server
        cluster.primaries.should include server
    end

    context "when talking to a replica set node" do

      context "that is not configured" do
        before do
          connection.pending_replies << Hash[
            "ismaster" => false,
            "secondary" => false,
            "info" => "can't get local.system.replset config from self or any seed (EMPTYCONFIG)",
            "isreplicaset" => true,
            "maxBsonObjectSize" => 16777216,
            "ok" => 1.0
          ]
        end
        it "returns nothing" do
          cluster.sync_server(server).should be_empty
      context "that is being initiated" do
        before do
          connection.pending_replies << Hash[
            "ismaster" => false,
            "secondary" => false,
            "info" => "Received replSetInitiate - should come online shortly.",
            "isreplicaset" => true,
            "maxBsonObjectSize" => 16777216,
            "ok" => 1.0
          ]
        end

        it "raises a connection failure exception" do
          cluster.sync_server(server).should be_empty
      context "that is ready but not elected" do
        before do
          connection.pending_replies << Hash[
            "setName" => "3fef4842b608",
            "ismaster" => false,
            "secondary" => false,
            "hosts" => ["localhost:61085", "localhost:61086", "localhost:61084"],
            "primary" => "localhost:61084",
            "me" => "localhost:61085",
            "maxBsonObjectSize" => 16777216,
            "ok" => 1.0
          ]
        end

        it "raises no exception" do
          lambda do
            cluster.sync_server server
          end.should_not raise_exception
        end
        it "adds the server to the list" do
          cluster.sync_server server
          cluster.servers.should include server
        it "returns all other known hosts" do
          cluster.sync_server(server).should =~ ["localhost:61085", "localhost:61086", "localhost:61084"]
      context "that is ready" do
        before do
          connection.pending_replies << Hash[
            "setName" => "3ff029114780",
            "ismaster" => true,
            "secondary" => false,
            "hosts" => ["localhost:59246", "localhost:59248", "localhost:59247"],
            "primary" => "localhost:59246",
            "me" => "localhost:59246",
            "maxBsonObjectSize" => 16777216,
            "ok" => 1.0
          ]
        end
        it "adds the node to the master set" do
          cluster.sync_server server
          cluster.primaries.should include server
        it "returns all other known hosts" do
          cluster.sync_server(server).should =~ ["localhost:59246", "localhost:59248", "localhost:59247"]

  describe "#socket_for" do
    let(:cluster) do
      Moped::Cluster.new ""
    let(:server) do
      Moped::Server.new("localhost:27017").tap do |server|
        server.stub(socket: socket)
    let(:socket) do
      Moped::Socket.new("127.0.0.1", 27017).tap do |socket|
        socket.connect
    context "when socket is dead" do
      let(:dead_server) do
        Moped::Server.allocate.tap do |server|
          server.stub(socket: dead_socket)
      let(:dead_socket) do
        Moped::Socket.new("127.0.0.1", 27017).tap do |socket|
          socket.stub(:alive? => false)
        primaries = [server, dead_server]
        primaries.stub(:sample).and_return(dead_server, server)
        cluster.stub(primaries: primaries)
      it "removes the socket" do
        cluster.should_receive(:remove).with(dead_server)
        cluster.socket_for :write
      it "returns the living socket" do
        cluster.socket_for(:write).should eq socket
    context "when mode is write" do
        server.primary = true
      context "and the cluster is not synced" do
        it "syncs the cluster" do
          cluster.should_receive(:sync) do
            cluster.servers << server
          cluster.socket_for :write
        it "returns the socket" do
          cluster.stub(:sync) { cluster.servers << server }
          cluster.socket_for(:write).should eq socket
        it "applies the cached authentication" do
          cluster.stub(:sync) { cluster.servers << server }
          socket.should_receive(:apply_auth).with(cluster.auth)
          cluster.socket_for(:write)
        end
      context "and the cluster is synced" do
        before do
          cluster.servers << server
        it "does not re-sync the cluster" do
          cluster.should_receive(:sync).never
          cluster.socket_for :write
        it "returns the socket" do
          cluster.socket_for(:write).should eq socket
        it "applies the cached authentication" do
          socket.should_receive(:apply_auth).with(cluster.auth)
          cluster.socket_for(:write)
    context "when mode is read" do
      context "and the cluster is not synced" do
        before do
          server.primary = true
        end
        it "syncs the cluster" do
          cluster.should_receive(:sync) do
            cluster.servers << server
          end
          cluster.socket_for :read
        it "applies the cached authentication" do
          cluster.stub(:sync) { cluster.servers << server }
          socket.should_receive(:apply_auth).with(cluster.auth)
          cluster.socket_for(:read)
      context "and the cluster is synced" do
        context "and no secondaries are found" do
          before do
            server.primary = true
            cluster.servers << server
          end

          it "returns the master connection" do
            cluster.socket_for(:read).should eq socket
          end
          it "applies the cached authentication" do
            socket.should_receive(:apply_auth).with(cluster.auth)
            cluster.socket_for(:read)
          end
        context "and a slave is found" do
          it "returns a random slave connection" do
            secondaries = [server]
            cluster.stub(secondaries: secondaries)
            secondaries.should_receive(:sample).and_return(server)
            cluster.socket_for(:read).should eq socket
          end

          it "applies the cached authentication" do
            cluster.stub(secondaries: [server])
            socket.should_receive(:apply_auth).with(cluster.auth)
            cluster.socket_for(:read)
          end
  describe "#login" do
    let(:cluster) do
      Moped::Cluster.allocate
    end
    it "adds the credentials to the auth cache" do
      cluster.login("admin", "username", "password")
      cluster.auth.should eq("admin" => ["username", "password"])
  describe "#logout" do
    let(:cluster) do
      Moped::Cluster.allocate
      cluster.login("admin", "username", "password")
    it "removes the stored credentials" do
      cluster.logout :admin
      cluster.auth.should be_empty

    mock(Moped::Session)
  end

  let(:database) do
    mock(Moped::Database, session: session, name: "moped")
  end

  let(:collection) do
    described_class.new database, :users
  describe "#initialize" do

    it "stores the database" do
      collection.database.should eq database
    end

    it "stores the collection name" do
      collection.name.should eq :users
    end
  end

  describe "#indexes" do
    it "returns a new indexes instance" do
      collection.indexes.should be_an_instance_of Moped::Indexes
    end
  end

    it "drops the collection" do
      database.should_receive(:command).with(drop: :users)
      collection.drop
    end
  end

  describe "#find" do

    let(:selector) do
      Hash[ a: 1 ]
    end

    let(:query) do
      mock(Moped::Query)
    end

    it "returns a new Query" do
      Moped::Query.should_receive(:new).
        with(collection, selector).and_return(query)
      collection.find(selector).should eq query
    it "defaults to an empty selector" do
      Moped::Query.should_receive(:new).
        with(collection, {}).and_return(query)
      collection.find.should eq query

    before do
      session.should_receive(:with, :consistency => :strong).and_yield(session)
      session.stub safe?: false
    end

    context "when passed a single document" do

      it "inserts the document" do
        session.should_receive(:execute).with do |insert|
          insert.documents.should eq [{a: 1}]
        end
        collection.insert(a: 1)
      end
    context "when passed multiple documents" do
      it "inserts the documents" do
        session.should_receive(:execute).with do |insert|
          insert.documents.should eq [{a: 1}, {b: 2}]
        end
        collection.insert([{a: 1}, {b: 2}])
      end
require "spec_helper"

describe Moped::Cursor do
  let(:session) { mock Moped::Session }
  let(:query_operation) { Moped::Protocol::Query.allocate }
  let(:cursor) { Moped::Cursor.new(session, query_operation) }

  describe "#initialize" do
    it "stores the session" do
      cursor.session.should eq session
    end

    it "stores a copy of the query operation" do
      query_operation.should_receive(:dup).and_return(query_operation)
      cursor.query_op.should eq query_operation
    end

    describe "the get_more operation" do
      it "inherits the query's database" do
        cursor.get_more_op.database.should eq query_operation.database
      end

      it "inherits the query's collection" do
        cursor.get_more_op.collection.should eq query_operation.collection
      end

      it "inherits the query's limit" do
        cursor.get_more_op.limit.should eq query_operation.limit
      end
    end
  end

  describe "#more?" do
    context "when get more operation's cursor id is 0" do
      it "returns false" do
        cursor.get_more_op.cursor_id = 0
        cursor.more?.should be_false
      end
    end
    context "when get more operation's cursor id is not 0" do
      it "returns true" do
        cursor.get_more_op.cursor_id = 123
        cursor.more?.should be_true
      end
    end
  end

  describe "#limited?" do
    context "when original query's limit is greater than 0" do
      before do
        query_operation.limit = 20
      end

      it "returns true" do
        cursor.should be_limited
      end
    end

    context "when original query's limit is not greater than 0" do
      before do
        query_operation.limit = 0
      end

      it "returns true" do
        cursor.should_not be_limited
      end
    end
  end

  describe "#query" do
    let(:reply) do
      Moped::Protocol::Reply.allocate.tap do |reply|
        reply.cursor_id = 123
        reply.count = 1
        reply.documents = [{"a" => 1}]
      end
    end

    before do
      session.stub(query: reply)
    end

    context "when query is limited" do
      before do
        query_operation.limit = 21
        cursor.query query_operation
      end

      it "updates the more operation's limit" do
        cursor.get_more_op.limit.should eq 20
      end

      it "sets the kill cursor operation's cursor id" do
        cursor.kill_cursor_op.cursor_ids.should eq [reply.cursor_id]
      end

      it "sets the more operation's cursor id" do
        cursor.get_more_op.cursor_id.should eq reply.cursor_id
      end
    end

    context "when query is limited" do
      before do
        query_operation.limit = 0
        cursor.query query_operation
      end

      it "does not update the more operation's limit" do
        cursor.get_more_op.limit.should eq query_operation.limit
      end

      it "sets the kill cursor operation's cursor id" do
        cursor.kill_cursor_op.cursor_ids.should eq [reply.cursor_id]
      end

      it "sets the more operation's cursor id" do
        cursor.get_more_op.cursor_id.should eq reply.cursor_id
      end
    end

    it "returns the documents" do
      cursor.query(query_operation).should eq reply.documents
    end
  end

  describe "#each" do

    context "when query returns all available documents" do
      let(:reply) do
        Moped::Protocol::Reply.allocate.tap do |reply|
          reply.cursor_id = 0
          reply.count = 21
          reply.documents = [{"a" => 1}]
        end
      end

      before do
        session.stub(query: reply)
      end

      it "yields each document" do
        results = []
        cursor.each { |doc| results << doc }
        results.should eq reply.documents
      end

      it "does not get more" do
        session.should_receive(:query).once
        cursor.each {}
      end

      it "does not kill the cursor" do
        cursor.should_receive(:kill).never
        cursor.each {}
      end
    end

    context "when query is unlimited" do
      let(:reply) do
        Moped::Protocol::Reply.allocate.tap do |reply|
          reply.cursor_id = 10
          reply.count = 10
          reply.documents = [{"a" => 1}]
        end
      end

      let(:get_more_reply) do
        Moped::Protocol::Reply.allocate.tap do |reply|
          reply.cursor_id = 0
          reply.count = 21
          reply.documents = [{"a" => 1}]
        end
      end

      before do
        session.stub(:query).and_return(reply, get_more_reply)
      end

      it "yields each document" do
        results = []
        cursor.each { |doc| results << doc }
        results.should eq reply.documents + get_more_reply.documents
      end

      it "gets more twice" do
        session.should_receive(:query).twice
        cursor.each {}
      end

      it "does not kill the cursor" do
        cursor.should_receive(:kill).never
        cursor.each {}
      end
    end

    context "when query is limited" do
      let(:reply) do
        Moped::Protocol::Reply.allocate.tap do |reply|
          reply.cursor_id = 10
          reply.count = 10
          reply.documents = [{"a" => 1}]
        end
      end

      let(:get_more_reply) do
        Moped::Protocol::Reply.allocate.tap do |reply|
          reply.cursor_id = 10
          reply.count = 10
          reply.documents = [{"a" => 1}]
        end
      end

      before do
        query_operation.limit = 20
        session.stub(:query).and_return(reply, get_more_reply)
        session.stub(:execute)
      end

      it "yields each document" do
        results = []
        cursor.each { |doc| results << doc }
        results.should eq reply.documents + get_more_reply.documents
      end

      it "gets more twice" do
        session.should_receive(:query).at_least(2)
        cursor.each {}
      end

      it "kills the cursor" do
        cursor.should_receive(:kill).once
        cursor.each {}
      end
    end

  end
end
require "spec_helper"

describe Moped::Database do

  let(:session) do
    Moped::Session.new ""
  end

  let(:database) do
    Moped::Database.new(session, :admin)
  end

  describe "#initialize" do

    it "stores the session" do
      database.session.should eq session
    end

    it "stores the database name" do
      database.name.should eq :admin
    end
  end

  describe "#command" do

    before do
      session.stub(:with).and_yield(session)
    end

    it "runs the given command against the master connection" do
      session.should_receive(:with, :consistency => :strong).
        and_yield(session)
      session.should_receive(:simple_query) do |query|
        query.full_collection_name.should eq "admin.$cmd"
        query.selector.should eq(ismaster: 1)

        { "ok" => 1.0 }
      end

      database.command ismaster: 1
    end

    context "when the command fails" do

      it "raises an exception" do
        session.stub(simple_query: { "ok" => 0.0 })

        expect {
          database.command ismaster: 1
        }.to raise_exception(Moped::Errors::OperationFailure)
      end
    end
  end

  describe "#drop" do

    it "drops the database" do
      database.should_receive(:command).with(dropDatabase: 1)
      database.drop
    end
  end

  describe "#[]" do

    it "returns a collection with that name" do
      Moped::Collection.should_receive(:new).with(database, :users)
      database[:users]
    end
  end

  describe "#login" do

    it "logs in to the database with the username and password" do
      session.cluster.should_receive(:login).with(:admin, "username", "password")
      database.login("username", "password")
    end
  end

  describe "#log out" do

    it "logs out from the database" do
      session.cluster.should_receive(:logout).with(:admin)
      database.logout
    end
  end
end
require "spec_helper"

describe Moped::Errors do

  describe "OperationFailure" do
    let(:command) do
      Moped::Protocol::Query.allocate
    end

    let(:error_details) do
      { "$err"=>"invalid query", "code"=>12580 }
    end

    let(:error) do
      described_class::OperationFailure.new(command, error_details)
    end

    describe "#initialize" do
      it "stores the command which generated the error" do
        error.command.should eq command
      end

      it "stores the details about the error" do
        error.details.should eq error_details
      end
    end

    describe "#message" do
      it "includes the command that generated the error" do
        error.message.should include command.inspect
      end

      context "when code is included in error details" do
        let(:error_details) do
          { "err" => "invalid query", "code" => 12580 }
        end

        it "includes the code" do
          error.message.should include error_details["code"].to_s
        end

        it "includes the error code reference site" do
          error.message.should include Moped::Errors::ERROR_REFERENCE
        end

        it "includes the error message" do
          error.message.should include error_details["err"].inspect
        end
      end

      context "when err is in the error details" do
        let(:error_details) do
          { "err" => "invalid query" }
        end

        it "includes the error message" do
          error.message.should include error_details["err"].inspect
        end
      end

      context "when $err is in the error details" do
        let(:error_details) do
          { "$err" => "not master" }
        end

        it "includes the error message" do
          error.message.should include error_details["$err"].inspect
        end
      end

      context "when errmsg is in the error details" do
        let(:error_details) do
          { "errmsg" => "invalid query" }
        end

        it "includes the error message" do
          error.message.should include error_details["errmsg"].inspect
        end
      end
    end

  end

  describe "QueryFailure" do
    it "is a kind of OperationFailure" do
      Moped::Errors::QueryFailure.ancestors.should \
        include Moped::Errors::OperationFailure
    end
  end

end
  let(:session) { Moped::Session.new ["127.0.0.1:27017"], database: "moped_test" }
  let(:indexes) do
    described_class.new(session.current_database, :users)
  end

  after do
    session.command(deleteIndexes: "users", index: "*")
  describe "#each" do
    before do
      session[:"system.indexes"].insert(ns: "moped_test.users", key: { name: 1 }, name: "name_1")
    end

    it "yields all indexes on the collection" do
      indexes.to_a.should eq \
        session[:"system.indexes"].find(ns: "moped_test.users").to_a
    end
  describe "#[]" do
    before do
      session[:"system.indexes"].insert(ns: "moped_test.users", key: { name: 1 }, name: "name_1")
    end

    it "returns the index with the provided key" do
      indexes[name: 1]["name"].should eq "name_1"
    end
    let(:key) do
      Hash["location.latlong" => "2d", "name" => 1, "age" => -1]
    end

    context "with no options" do
      it "creates an index with a generated name" do
        indexes.create(key)
        indexes[key]["name"].should eq "location.latlong_2d_name_1_age_-1"
      end
    end

    context "with a name provided" do
      it "creates an index with the provided name" do
        indexes.create(key, name: "custom_index_name")
        indexes[key]["name"].should eq "custom_index_name"
      end
    end

    context "with background: true" do
      it "creates an index" do
        indexes.create(key, background: true)
        indexes[key]["background"].should eq true
      end
    end

    context "with dropDups: true" do
      it "creates an index" do
        indexes.create(key, dropDups: true)
        indexes[key]["dropDups"].should eq true
    context "with unique: true" do
      it "creates an index" do
        indexes.create(key, unique: true)
        indexes[key]["unique"].should eq true

    context "with sparse: true" do
      it "creates an index" do
        indexes.create(key, sparse: true)
        indexes[key]["sparse"].should eq true
      end
    end

    context "with v: 0" do
      it "creates an index" do
        indexes.create(key, v: 0)
        indexes[key]["v"].should eq 0
      end
    end

    before do
      indexes.create name: 1
      indexes.create age: -1
    end

    context "with no key" do
      before do
        indexes.drop
      end

      it "drops all indexes for the collection" do
        indexes[name: 1].should be_nil
        indexes[age: -1].should be_nil
    context "with a key" do
      before do
        indexes.drop(name: 1)
      end

      it "drops the index that matches the key" do
      end

      it "does not drop other indexes" do
        indexes[age: -1].should_not be_nil
      end
    end

    context "with a key that doesn't exist" do
      it "returns false" do
        indexes.drop(other: 1).should be_false
require "spec_helper"

describe Moped::Logging do
  let(:config) do
    Module.new { extend Moped::Logging }
  end
  let(:logger) { mock(Logger) }

  describe ".rails_logger" do
    context "when Rails is present" do
      let(:rails) { Class.new }

      before do
        Object.const_set :Rails, rails
      end

      after do
        Object.send(:remove_const, :Rails)
      end

      context "and it defines logger" do
        before do
          rails.stub(logger: logger)
        end

        it "returns the logger" do
          config.rails_logger.should eq logger
        end
      end

      context "but does not define logger" do
        it "returns false" do
          config.rails_logger.should be_false
        end
      end

    end
  end

  describe ".default_logger" do
    it "returns a new logger instance" do
      config.default_logger.should be_a_kind_of Logger
    end

    it "sets the log level to info" do
      config.default_logger.level.should eq Logger::INFO
    end
  end

  describe ".logger" do
    context "when a rails logger is available" do
      before do
        config.stub(rails_logger: logger)
      end

      it "returns the rails logger" do
        config.logger.should eq logger
      end
    end

    context "when a rails logger is not available" do
      before do
        config.stub(rails_logger: nil)
        config.stub(default_logger: logger)
      end

      it "returns the default logger" do
        config.logger.should eq logger
      end
    end

    context "when the logger is set to nil" do
      before do
        config.logger = nil
      end

      it "returns nil" do
        config.logger.should be_nil
      end
    end
  end

end
  let(:session) do
    mock(Moped::Session)
  end

  let(:database) do
    mock(
      Moped::Database,
      name: "moped",
      session: session
    )
  end

  let(:collection) do
    mock(
      Moped::Collection,
      database: database,
      name: "users"
    )
  end

  let(:selector) do
    Hash[ a: 1 ]
  end

  let(:query) do
    described_class.new collection, selector
  end

  describe "#initialize" do

    it "stores the collection" do
      query.collection.should eq collection
    it "stores the selector" do
      query.selector.should eq selector
  end
  describe "#limit" do
    it "sets the query operation's limit field" do
      query.limit(5)
      query.operation.limit.should eq 5
    it "returns the query" do
      query.limit(5).should eql query
  end
  describe "#skip" do
    it "sets the query operation's skip field" do
      query.skip(5)
      query.operation.skip.should eq 5
    it "returns the query" do
      query.skip(5).should eql query
  end
  describe "#select" do
    it "sets the query operation's fields" do
      query.select(a: 1)
      query.operation.fields.should eq(a: 1)
    it "returns the query" do
      query.select(a: 1).should eql query
    end
  end
  describe "#sort" do
    context "when called for the first time" do
      it "updates the selector to mongo's advanced selector" do
        query.sort(a: 1)
        query.operation.selector.should eq(
          "$query" => selector,
          "$orderby" => { a: 1 }
        )
    context "when called again" do
      it "changes the $orderby" do
        query.sort(a: 1)
        query.sort(a: 2)
        query.operation.selector.should eq(
          "$query" => selector,
          "$orderby" => { a: 2 }
        )
    it "returns the query" do
      query.sort(a: 1).should eql query
    end
  end
  describe "#explain" do
    before do
      session.should_receive(:simple_query).with(query.operation)
    end
    context "when a sort exists" do
      before do
        query.sort(_id: 1)
      it "updates to a mongo advanced selector" do
        query.explain
        query.operation.selector.should eq(
          "$query" => selector,
          "$explain" => true,
          "$orderby" => { _id: 1 }
        )
      end
    end
    context "when no sort exists" do
      it "updates to a mongo advanced selector" do
        query.explain
        query.operation.selector.should eq(
          "$query" => selector,
          "$explain" => true,
          "$orderby" => {}
        )
  end
  describe "#one" do
    it "executes a simple query" do
      session.should_receive(:simple_query).with(query.operation)
      query.one
  end
  describe "#distinct" do
    it "executes a distinct command" do
      database.should_receive(:command).with(
        distinct: collection.name,
        key: "name",
        query: selector
      ).and_return("values" => [ "durran", "bernerd" ])
      query.distinct(:name)
  end
  describe "#count" do
    it "executes a count command" do
      database.should_receive(:command).with(
        count: collection.name,
        query: selector
      ).and_return("n" => 4)
      query.count
    it "returns the count" do
      database.stub(command: { "n" => 4 })
      query.count.should eq 4
  describe "#update" do

    let(:change) do
      Hash[ a: 1 ]
    it "updates the record matching selector with change" do
      session.should_receive(:with, :consistency => :strong).
        and_yield(session)
      session.should_receive(:execute).with do |update|
        update.flags.should eq []
        update.selector.should eq query.operation.selector
        update.update.should eq change
      end

      query.update change
    end
  end
  describe "#update_all" do
    let(:change) do
      Hash[ a: 1 ]
    end
    it "updates all records matching selector with change" do
      query.should_receive(:update).with(change, [:multi])
      query.update_all change
  describe "#upsert" do

    let(:change) do
      Hash[ a: 1 ]
    it "upserts the record matching selector with change" do
      query.should_receive(:update).with(change, [:upsert])
      query.upsert change
    end
  describe "#remove" do
    it "removes the first matching document" do
      session.should_receive(:with, :consistency => :strong).
        and_yield(session)
      session.should_receive(:execute).with do |delete|
        delete.flags.should eq [:remove_first]
        delete.selector.should eq query.operation.selector
      end

      query.remove
    end
  describe "#remove_all" do
    it "removes all matching documents" do
      session.should_receive(:with, :consistency => :strong).
        and_yield(session)
      session.should_receive(:execute).with do |delete|
        delete.flags.should eq []
        delete.selector.should eq query.operation.selector
      end
      query.remove_all
  end
  describe "#each" do
      session.should_receive(:with).
        with(retain_socket: true).and_return(session)
    it "creates a new cursor" do
      cursor = mock(Moped::Cursor, next: nil)
      Moped::Cursor.should_receive(:new).
        with(session, query.operation).and_return(cursor)
      query.each
    end
    it "yields all documents in the cursor" do
      cursor = Moped::Cursor.allocate
      cursor.stub(:to_enum).and_return([1, 2].to_enum)
      Moped::Cursor.stub(new: cursor)
      query.to_a.should eq [1, 2]
    it "returns an enumerator" do
      cursor = mock(Moped::Cursor)
      Moped::Cursor.stub(new: cursor)
      query.each.should be_a Enumerator
require "spec_helper"

describe Moped::Server do

  describe "#initialize" do
    let(:server) do
      described_class.new("localhost:123")
    end

    it "stores the original address" do
      server.address.should eq "localhost:123"
    end

    it "stores the resolved address" do
      server.resolved_address.should eql "127.0.0.1:123"
    end

    it "stores the resolved ip" do
      server.ip_address.should eq "127.0.0.1"
    end

    it "stores the port" do
      server.port.should eq 123
    end
  end

  describe "==" do
    context "when ip and port are the same" do
      it "returns true" do
        described_class.new("127.0.0.1:999").should eq \
          described_class.new("localhost:999")
      end
    end

    context "when ip and port are different" do
      it "returns false" do
        described_class.new("127.0.0.1:1000").should_not eq \
          described_class.new("localhost:999")
      end
    end

    context "when other is not a server" do
      it "returns false" do
        described_class.new("127.0.0.1:999").should_not eq 1
      end
    end
  end

  context "when added to a set" do
    let(:set) { Set.new }

    context "and ip and port are the same" do
      it "does not add both servers" do
        set << described_class.new("127.0.0.1:1000")
        set << described_class.new("127.0.0.1:1000")

        set.length.should eq 1
      end

      context "and the original address is different" do
        it "does not add both servers" do
          set << described_class.new("localhost:1000")
          set << described_class.new("127.0.0.1:1000")

          set.length.should eq 1
        end
      end
    end

    context "and ip and port are different" do
      it "adds both servers" do
        set << described_class.new("127.0.0.1:1000")
        set << described_class.new("127.0.0.1:2000")

        set.length.should eq 2
      end
    end
  end

end

  let(:seeds) do
    "127.0.0.1:27017"
  end

  let(:options) do
    Hash[database: "test", safe: true, consistency: :eventual]
  end

    described_class.new seeds, options
  end

  describe "#initialize" do

    it "stores the options provided" do
      session.options.should eq(options)
    end

    it "stores the cluster" do
      session.cluster.should be_a(Moped::Cluster)
    end
  end

  describe "#current_database" do

    context "when no database option has been set" do

      let(:session) do
        described_class.new seeds, {}
      end

      it "raises an exception" do
        expect { session.current_database }.to raise_exception
      end
    end

    context "when a database option is set" do

      let(:database) do
        stub
      end

      before do
        Moped::Database.should_receive(:new).
          with(session, options[:database]).and_return(database)
      end

      it "returns the database from the options" do
        session.current_database.should eq(database)
      end

      it "memoizes the database" do
        database = session.current_database
        session.current_database.should equal(database)
      end
    end
  end

  describe "#safe?" do

    context "when :safe is not present" do

      before do
        session.options.delete(:safe)
      end

      it "returns false" do
        session.should_not be_safe
      end
    end

    context "when :safe is present but false" do

      before do
        session.options[:safe] = false
      end

      it "returns false" do
        session.should_not be_safe
      end
    end

    context "when :safe is true" do

      before do
        session.options[:safe] = true
      end

      it "returns true" do
        session.should be_safe
      end
    end

    context "when :safe is a hash" do

      before do
        session.options[:safe] = { fsync: true }
      end

      it "returns true" do
        session.should be_safe
      end
    end

    it "sets the :database option" do
      session.use :admin
      session.options[:database].should eq(:admin)
    end

    context "when there is not already a current database" do

      it "sets the current database" do
        session.should_receive(:set_current_database).with(:admin)
        session.use :admin
      end

    let(:new_options) do
      Hash[database: "test-2"]
    end


      it "yields a session" do
        session.with(new_options) do |new_session|
          new_session.should be_a Moped::Session
        end
      it "yields a new session" do
        session.with(new_options) do |new_session|
          new_session.should_not eql session
      it "returns the result of the block" do
        session.with(new_options) { false }.should eq false
      end

      it "merges the old and new session's options" do
        session.with(new_options) do |new_session|
          new_session.options.should eq options.merge(new_options)
        end
      end

      it "does not change the original session's options" do
        original_options = options.dup
        session.with(new_options) do |new_session|
          session.options.should eql original_options
        end
      end

      it "unmemoizes the current database" do
        db = session.current_database
        session.with(new_options) do |new_session|
          new_session.current_database.should_not eql db

      it "returns a session" do
        session.with(new_options).should be_a Moped::Session
      end

      it "returns a new session" do
        session.with(new_options).should_not eql session
      end

      it "merges the old and new session's options" do
        session.with(new_options).options.should eq options.merge(new_options)
      end

      it "does not change the original session's options" do
        original_options = options.dup
        session.with(new_options)
        session.options.should eql original_options
      end
    end
  end

  describe "#new" do

    let(:new_options) do
      Hash[database: "test-2"]
    end

    let(:new_session) do
      described_class.new seeds, options
    end

    before do
      new_session.cluster.stub(:reconnect)
    end

    it "delegates to #with" do
      session.should_receive(:with).with(new_options).and_return(new_session)
      session.new(new_options)
    end

    it "instructs the cluster to reconnect" do
      session.stub(with: new_session)
      new_session.cluster.should_receive(:reconnect)
      session.new(new_options)
    end

    context "when called with a block" do

      it "yields the new session" do
        session.stub(with: new_session)
        session.new(new_options) do |session|
          session.should eql new_session
        end
    end
    context "when called without a block" do

      it "returns the new session" do
        session.stub(with: new_session)
        session.new(new_options).should eql new_session

    it "delegates to the current database" do
      database = mock(Moped::Database)
      session.should_receive(:current_database).and_return(database)
      database.should_receive(:drop)
      session.drop
    end
  end

  describe "#command" do

    let(:command) do
      Hash[ismaster: 1]
    end

    it "delegates to the current database" do
      database = mock(Moped::Database)
      session.should_receive(:current_database).and_return(database)
      database.should_receive(:command).with(command)
      session.command command
    end
  end

  describe "#login" do

    it "delegates to the current database" do
      database = mock(Moped::Database)
      session.should_receive(:current_database).and_return(database)
      database.should_receive(:login).with("username", "password")
      session.login("username", "password")
    end
  end

  describe "#logout" do

    it "delegates to the current database" do
      database = mock(Moped::Database)
      session.should_receive(:current_database).and_return(database)
      database.should_receive(:logout)
      session.logout
    end
  end

  describe "#socket_for" do

    it "delegates to the cluster" do
      session.cluster.should_receive(:socket_for).with(:read)
      session.send(:socket_for, :read)
    end

    context "when retain socket option is set" do

      before do
        session.options[:retain_socket] = true
      end

      it "only aquires the socket once" do
        session.cluster.should_receive(:socket_for).
          with(:read).once.and_return(mock(Moped::Socket))

        session.send(:socket_for, :read)
        session.send(:socket_for, :read)
      end
    end
  end

  describe "#simple_query" do

    let(:query) do
      Moped::Protocol::Query.allocate
    end

    let(:socket) do
      mock(Moped::Socket)
    end

    let(:reply) do
      Moped::Protocol::Reply.allocate.tap do |reply|
        reply.documents = [{a: 1}]

    before do
      session.stub(socket_for: socket)
      session.stub(query: reply)
    end

    it "limits the query" do
      session.should_receive(:query) do |query|
        query.limit.should eq(-1)
        reply
      end

      session.simple_query(query)
    end

    it "returns the document" do
      session.simple_query(query).should eq(a: 1)
    end
  describe "#query" do

    let(:query) do
      Moped::Protocol::Query.allocate
    end

    let(:socket) do
      mock(Moped::Socket)
    end

    let(:reply) do
      Moped::Protocol::Reply.allocate.tap do |reply|
        reply.documents = [{a: 1}]
      end
    end

    before do
      session.stub(socket_for: socket)
      socket.stub(:execute).and_return(reply)
    end

    context "when consistency is strong" do

      before do
        session.options[:consistency] = :strong
      end

      it "queries the master node" do
        session.should_receive(:socket_for).with(:write).
          and_return(socket)
        session.query(query)
      end
    end

    context "when consistency is eventual" do

      before do
        session.options[:consistency] = :eventual
      end

      it "queries a slave node" do
        session.should_receive(:socket_for).with(:read).
          and_return(socket)
        session.query(query)
      end

      context "and query accepts flags" do

        it "sets slave_ok on the query flags" do
          session.stub(socket_for: socket)
          socket.should_receive(:execute) do |query|
            query.flags.should include :slave_ok
          end

          session.query(query)
        end
      end

      context "and query does not accept flags" do

        let(:query) do
          Moped::Protocol::GetMore.allocate
        end

        it "doesn't try to set flags" do
          session.stub(socket_for: socket)
          expect { session.query(query) }.not_to raise_exception
        end
      end
    end

    context "when reply has :query_failure flag" do

      before do
        reply.flags = [:query_failure]
      end

      it "raises a QueryFailure exception" do
        expect {
          session.query(query)
        }.to raise_exception(Moped::Errors::QueryFailure)
  describe "#execute" do

    let(:operation) do
      Moped::Protocol::Insert.allocate
    end

    let(:socket) do
      mock(Moped::Socket)
    end

    context "when session is not in safe mode" do

      before do
        session.options[:safe] = false
      end

      context "when consistency is strong" do

        before do
          session.options[:consistency] = :strong
        end

        it "executes the operation on the master node" do
          session.should_receive(:socket_for).with(:write).
            and_return(socket)
          socket.should_receive(:execute).with(operation)

          session.execute(operation)
        end
      end

      context "when consistency is eventual" do

        before do
          session.options[:consistency] = :eventual
        end

        it "executes the operation on a slave node" do
          session.should_receive(:socket_for).with(:read).
            and_return(socket)
          socket.should_receive(:execute).with(operation)

          session.execute(operation)
        end
      end
    end

    context "when session is in safe mode" do

      let(:reply) do
        Moped::Protocol::Reply.allocate.tap do |reply|
          reply.documents = [{a: 1}]
        end
      end

      before do
        session.options[:safe] = { w: 2 }
      end

      context "when the operation fails" do

        let(:reply) do
          Moped::Protocol::Reply.allocate.tap do |reply|
            reply.documents = [{
              "err"=>"document to insert can't have $ fields",
              "code"=>13511,
              "n"=>0,
              "connectionId"=>894,
              "ok"=>1.0
            }]
          end
        end

        it "raises an OperationFailure exception" do
          session.stub(socket_for: socket)
          socket.stub(execute: reply)

          expect {
            session.execute(operation)
          }.to raise_exception(Moped::Errors::OperationFailure)
        end
      end

      context "when consistency is strong" do

        before do
          session.options[:consistency] = :strong
        end

        it "executes the operation on the master node" do
          session.should_receive(:socket_for).with(:write).
            and_return(socket)

          socket.should_receive(:execute) do |op, query|
            op.should eq operation
            query.selector.should eq(getlasterror: 1, w: 2)
            reply
          end

          session.execute(operation)
        end
      end

      context "when consistency is eventual" do

        before do
          session.options[:consistency] = :eventual
        end

        it "executes the operation on a slave node" do
          session.should_receive(:socket_for).with(:read).
            and_return(socket)

          socket.should_receive(:execute) do |op, query|
            op.should eq operation
            query.selector.should eq(getlasterror: 1, w: 2)
            reply
          end

          session.execute(operation)
        end
      end
    end
  end
require "spec_helper"

describe Moped::Socket do

  let!(:server) do
    TCPServer.new "127.0.0.1", 0
  end

  let(:socket) do
    described_class.new "127.0.0.1", server.addr[1]
  end

  let(:connection) {
    socket.connection
  }

  before do
    socket.connect
  end

  after do
    connection.close if connection && !connection.closed?
    server.close unless server.closed?
  end

  describe "#initialize" do
    it "stores the host of the server" do
      socket.host.should eq "127.0.0.1"
    end

    it "stores the port of the server" do
      socket.port.should eq server.addr[1]
    end

    it "connects to the server" do
      socket.connection.should_not be_closed
    end
  end

  describe "#connect" do
    context "when node is not running" do
      let(:bogus_port) do
        server = TCPServer.new("127.0.0.1", 0)
        server.addr[1].tap do
          server.close
        end
      end

      let(:socket) do
        described_class.new "127.0.0.1", bogus_port
      end

      it "returns false" do
        socket.connect.should be_false
      end
    end

    context "when connection times out" do
      if RUBY_PLATFORM == "java"
        let(:timeout_server) do
          java.net.ServerSocket.new(0, 1)
        end

        let(:timeout_port) do
          timeout_server.getLocalPort
        end
      else
        let(:timeout_server) do
          TCPServer.new("127.0.0.1", 0).tap do |server|
            server.listen(1)
          end
        end

        let(:timeout_port) do
          timeout_server.addr[1]
        end
      end

      let(:timeout_socket) do
        described_class.new "127.0.0.1", timeout_port
      end

      before do
        sockaddr = Socket.pack_sockaddr_in(timeout_port, '127.0.0.1')

        5.times do # flood the server socket
          ::Socket.new(::Socket::AF_INET, ::Socket::SOCK_STREAM, 0).connect_nonblock(sockaddr) rescue nil
        end
      end

      after do
        timeout_server.close unless timeout_server.closed?
      end

      it "returns false" do
        timeout_socket.connect.should be_false
      end
    end
  end

  describe "#alive?" do
    context "when not connected" do
      let(:socket) do
        described_class.new("127.0.0.1", 99999).tap do |socket|
          socket.stub(:connect)
        end
      end

      it "should be false" do
        socket.should_not be_alive
      end
    end

    context "when connected but server goes away" do
      before do
        remote = server.accept
        remote.shutdown
        # Give the socket time to be notified
        sleep 0.1
      end

      it "should be false" do
        socket.should_not be_alive
      end
    end

    context "when connected but server goes away" do
      before do
        server.close
        # Give the socket time to be notified
        sleep 0.1
      end

      it "should be false" do
        socket.should_not be_alive
      end
    end

    context "when connect is explicitly closed" do
      before do
        socket.close
      end

      it "should be false" do
        socket.should_not be_alive
      end
    end

    context "when connected and server is open" do
      it "should be true" do
        socket.should be_alive
      end
    end
  end

  describe "#execute" do
    let(:query) { Moped::Protocol::Query.new(:moped, :test, {}) }

    context "when competing threads attempt to query" do
      let(:messages) do
        10.times.map do |i|
          Moped::Protocol::Insert.new(:moped, :test, {}).tap do |query|
            query.stub(request_id: 123)
          end
        end
      end

      it "never issues a partial write" do
        socket

        threads = 10.times.map do |i|
          Thread.new do
            Thread.current.abort_on_exception = true
            socket.execute messages[i]
          end
        end

        threads.each(&:join)

        sock = server.accept
        data = sock.read messages.join.length

        messages.each do |message|
          fail "server received partial write" unless data.include? message.serialize
        end
      end
    end
  end

  describe "#parse_reply" do
    let(:raw) do
      Moped::Protocol::Reply.allocate.tap do |reply|
        reply.request_id = 1
        reply.response_to = 1
        reply.op_code = 1
        reply.flags = [:await_capable]
        reply.offset = 4
        reply.count = 1
        reply.documents = [{"name" => "John"}]
      end.serialize
    end

    let(:reply) do
      socket.parse_reply(raw.length, raw[4..-1])
    end

    it "sets the length" do
      reply.length.should eq raw.length
    end

    it "sets the response_to" do
      reply.response_to.should eq 1
    end

    it "sets the request id" do
      reply.request_id.should eq 1
    end

    it "sets the flags" do
      reply.flags.should eq [:await_capable]
    end

    it "sets the offset" do
      reply.offset.should eq 4
    end

    it "sets the count" do
      reply.count.should eq 1
    end

    it "sets the documents" do
      reply.documents.should eq [{"name" => "John"}]
    end
  end

  describe "#close" do
    let(:exception) { RuntimeError.new }
    let(:callback) { stub }

    it "closes the connection" do
      connection.should_receive(:close).at_least(1)
      socket.close
    end

    it "marks the socket as dead" do
      socket.close
      socket.should_not be_alive
    end
  end

  describe "#login" do

    let(:connection) do
      Support::MockConnection.new
    end

    before do
      socket.stub(connection: connection)
    end

    context "when authentication is successful" do
      before do
        # getnonce
        connection.pending_replies << Hash["nonce" => "123", "ok" => 1]
        # authenticate
        connection.pending_replies << Hash["ok" => 1]
      end

      it "returns true" do
        socket.login("admin", "username", "password").should be_true
      end

      it "adds the credentials to the auth cache" do
        socket.login(:admin, "username", "password")
        socket.auth.should eq("admin" => ["username", "password"])
      end
    end

    context "when a nonce fails to generate" do
      before do
        # getnonce
        connection.pending_replies << Hash["ok" => 0]
      end

      it "raises an operation failure" do
        lambda do
          socket.login(:admin, "username", "password")
        end.should raise_exception(Moped::Errors::OperationFailure)
      end

      it "does not add the credentials to the auth cache" do
        socket.login(:admin, "username", "password") rescue nil
        socket.auth.should be_empty
      end
    end

    context "when authentication fails" do
      before do
        # getnonce
        connection.pending_replies << Hash["nonce" => "123", "ok" => 1]
        # authenticate
        connection.pending_replies << Hash["ok" => 0]
      end

      it "raises an operation failure" do
        lambda do
          socket.login(:admin, "username", "password")
        end.should raise_exception(Moped::Errors::OperationFailure)
      end

      it "does not add the credentials to the auth cache" do
        socket.login(:admin, "username", "password") rescue nil
        socket.auth.should be_empty
      end
    end

  end

  describe "#logout" do

    let(:connection) do
      Support::MockConnection.new
    end

    before do
      socket.stub(connection: connection)
      socket.auth["admin"] = ["username", "password"]
    end

    context "when logout is successful" do
      before do
        connection.pending_replies << Hash["ok" => 1]
      end

      it "removes the stored credentials" do
        socket.logout :admin
        socket.auth.should be_empty
      end
    end

    context "when logout is unsuccessful" do
      before do
        connection.pending_replies << Hash["ok" => 0]
      end

      it "does not remove the stored credentials" do
        socket.logout :admin rescue nil
        socket.auth.should_not be_empty
      end

      it "raises an operation failure" do
        lambda do
          socket.logout :admin
        end.should raise_exception(Moped::Errors::OperationFailure)
      end
    end

  end

  describe "#apply_auth" do
    context "when the socket is unauthenticated" do
      it "logs in with each credential provided" do
        socket.should_receive(:login).with("admin", "username", "password")
        socket.should_receive(:login).with("test", "username", "password")

        socket.apply_auth(
          "admin" => ["username", "password"],
          "test" => ["username", "password"]
        )
      end
    end

    context "when the socket is authenticated" do
      before do
        socket.auth["admin"] = ["username", "password"]
      end

      context "and a credential is unchanged" do
        it "does nothing" do
          socket.should_not_receive(:login)
          socket.apply_auth("admin" => ["username", "password"])
        end
      end

      context "and a credential changes" do
        it "logs in with the new credentials" do
          socket.should_receive(:login).with("admin", "newuser", "password")
          socket.apply_auth("admin" => ["newuser", "password"])
        end
      end

      context "and a credential is removed" do
        it "logs out from the database" do
          socket.should_receive(:logout).with("admin")
          socket.apply_auth({})
        end
      end

      context "and a credential is added" do
        it "logs in with the added credentials" do
          socket.should_receive(:login).with("test", "username", "password")
          socket.apply_auth(
            "admin" => ["username", "password"],
            "test" => ["username", "password"]
          )
        end
      end
    end
  end

  describe "instrument" do

    context "when a logger is configured in debug mode" do
      before do
        Moped.stub(logger: mock(Logger, debug?: true))
      end

      it "logs the operations" do
        socket.should_receive(:log_operations).once
        socket.instrument([]) {}
      end
    end

    context "when a logger is configured but not in debug level" do
      before do
        Moped.stub(logger: mock(Logger, debug?: false))
      end

      it "does not log the operations" do
        socket.should_receive(:log_operations).never
        socket.instrument([]) {}
      end
    end

    context "when no logger is configured" do
      before do
        Moped.stub(logger: nil)
      end

      it "does not log the operations" do
        socket.should_receive(:log_operations).never
        socket.instrument([]) {}
      end
    end

    context "when an error occurs" do
      before do
        Moped.stub(logger: mock(Logger, debug?: true))
      end

      it "does not log the operations" do
        socket.should_receive(:log_operations).never

        lambda do
          socket.instrument([]) { raise "inner error" }
        end.should raise_exception("inner error")
      end
    end

  end

end
require "spec_helper"
require "spec_helper"

describe "testing" do
  let(:cluster) { Moped::Cluster.new "", false }
  let(:socket) { Moped::Socket.new "", 99999 }
  let(:connection) { Support::MockConnection.new }

  before do
    socket.stub(connection: connection, alive?: true)
  end

  describe "#sync_socket" do
  end
end

__END__

# sequence for single node startup:
#
#   connection failure (node not up)

{"ismaster" => true, "maxBsonObjectSize" => 16777216, "ok" => 1.0}

# sequence for replica set startup pre-initialize:
#
#   connection failure

{"ismaster" => false, "secondary" => false, "info" => "can't get local.system.replset config from self or any seed (EMPTYCONFIG)", "isreplicaset" => true, "maxBsonObjectSize" => 16777216, "ok" => 1.0}

# sequence for replica set startup (master):
#
#   connection failure (node not up)

{"ismaster" => false, "secondary" => false, "info" => "can't get local.system.replset config from self or any seed (EMPTYCONFIG)", "isreplicaset" => true, "maxBsonObjectSize" => 16777216, "ok" => 1.0}

{"setName" => "3ff029114780", "ismaster" => false, "secondary" => true, "hosts" => ["localhost:59246", "localhost:59248", "localhost:59247"], "me" => "localhost:59246", "maxBsonObjectSize" => 16777216, "ok" => 1.0}

{"ismaster" => false, "secondary" => false, "info" => "Received replSetInitiate - should come online shortly.", "isreplicaset" => true, "maxBsonObjectSize" => 16777216, "ok" => 1.0}

{"setName" => "3ff029114780", "ismaster" => true, "secondary" => false, "hosts" => ["localhost:59246", "localhost:59248", "localhost:59247"], "primary" => "localhost:59246", "me" => "localhost:59246", "maxBsonObjectSize" => 16777216, "ok" => 1.0}

# sequence for replica set startup (secondary):
#
#   connection failure (node not up)

{"ismaster" => false, "secondary" => false, "info" => "can't get local.system.replset config from self or any seed (EMPTYCONFIG)", "isreplicaset" => true, "maxBsonObjectSize" => 16777216, "ok" => 1.0}

{"setName" => "3fef4842b608", "ismaster" => false, "secondary" => false, "hosts" => ["localhost:61085", "localhost:61086", "localhost:61084"], "me" => "localhost:61085", "maxBsonObjectSize" => 16777216, "ok" => 1.0}

{"setName" => "3fef4842b608", "ismaster" => false, "secondary" => false, "hosts" => ["localhost:61085", "localhost:61086", "localhost:61084"], "primary" => "localhost:61084", "me" => "localhost:61085", "maxBsonObjectSize" => 16777216, "ok" => 1.0}

{"setName" => "3fef4842b608", "ismaster" => false, "secondary" => true, "hosts" => ["localhost:61085", "localhost:61086", "localhost:61084"], "primary" => "localhost:61084", "me" => "localhost:61085", "maxBsonObjectSize" => 16777216, "ok" => 1.0}

__END__

describe Moped::Cluster do
  context "when connecting to a single seed" do
    context "and the seed is down"

    context "and that seed is primary" do
      it "finds the primary node"

      it "adds the secondary to the dynamic seeds"
      it "adds the arbiter to the dynamic seeds"
    end

    context "and that seed is secondary" do
      it "finds the secondary node"

      it "adds the primary to the dynamic seeds"
      it "adds the arbiter to the dynamic seeds"
    end

    context "and that seed is an arbiter" do
      it "adds the primary to the dynamic seeds"
      it "adds the secondary to the dynamic seeds"
    end
  end

  context "when connected to a single seed" do
    context "and that seed goes down" do
      it "is able to resync from discovered nodes"
    end
  end

  context "when connecting to a replica set" do
    context "and the replica set is not initiated"
    context "and the replica set is partially initiated"
    context "and there is no master node"
    context "and there is no secondary node"
  end

  context "when connected to a replica set" do
    context "and the primary node goes down" do
      it "issues reads to the secondary"
    end

    context "and the primary node changes"

    context "and the secondary node goes down" do
      it "issues inserts to the primary"
      it "issues reads to the primary"
    end
  end

end
require "bundler"
Bundler.require
require "support/mock_connection"
module Support
  class MockConnection

    attr_reader :pending_replies

    def initialize
      @connected = true
      @buffer = StringIO.new
      @pending_replies = []
    end

    def write(*args)

    end

    def read(*args)
      while documents = pending_replies.shift
        reply = Moped::Protocol::Reply.allocate.tap do |reply|
          documents = [documents] unless documents.is_a? Array
          reply.documents = documents
          reply.count = documents.length
        end

        reply.serialize(@buffer.string)
      end

      @buffer.read(*args)
    end

    def closed?
      !@connected
    end

    def close
      @connected = false
    end

  end
end
Fixed code: .env
## Features

* Automated replica set node discovery and failover.
* No C or Java extensions
* No external dependencies
* Simple, stable, public API.

### Unsupported Features

* GridFS
* Map/Reduce

These features are possible to implement, but outside the scope of Moped's
goals. Consider them perfect opportunities to write a companion gem!

<table><tbody>

<tr><th>new</th>
<td>Creates a new object id.</td></tr>

<tr><th>from_string</th>
<td>Creates a new object id from an object id string.
<br>
<code>Moped::BSON::ObjectId.from_string("4f8d8c66e5a4e45396000009")</code>
</td></tr>

<tr><th>from_time</th>
<td>Creates a new object id from a time.
<br>
<code>Moped::BSON::ObjectId.from_time(Time.new)</code>
</td></tr>

<tr><th>legal?</th>
<td>Validates an object id string.
<br>
<code>Moped::BSON::ObjectId.legal?("4f8d8c66e5a4e45396000009")</code>
</td></tr>

</tbody></table>

# Exceptions

Here's a list of the exceptions generated by Moped.

<table><tbody>

<tr><th>Moped::Errors::ConnectionFailure</th>
<td>Raised when a node cannot be reached or a connection is lost.
<br>
<strong>Note:</strong> this exception is only raised if Moped could not
reconnect, so you shouldn't attempt to rescue this.</td></tr>

<tr><th>Moped::Errors::OperationFailure</th>
<td>Raised when a command fails or is invalid, such as when an insert fails in
safe mode.</td></tr>

<tr><th>Moped::Errors::QueryFailure</th>
<td>Raised when an invalid query was sent to the database.</td></tr>

<tr><th>Moped::Errors::AuthenticationFailure</th>
<td>Raised when invalid credentials were passed to `session.login`.</td></tr>

<tr><th>Moped::Errors::SocketError</th>
<td>Not a real exception, but a module used to tag unhandled exceptions inside
of a node's networking code. Allows you to `rescue Moped::SocketError` which
preserving the real exception.</td></tr>

</tbody></table>

Other exceptions are possible while running commands, such as IO Errors around
failed connections. Moped tries to be smart about managing its connections,
such as checking if they're dead before executing a command; but those checks
aren't foolproof, and Moped is conservative about handling unexpected errors on
its connections. Namely, Moped will *not* retry a command if an unexpected
exception is raised. Why? Because it's impossible to know whether the command
was actually received by the remote Mongo instance, and without domain
knowledge it cannot be safely retried.

Take for example this case:

```ruby
session.with(safe: true)["users"].insert(name: "John")
```

It's entirely possible that the insert command will be sent to Mongo, but the
connection gets closed before we read the result for `getLastError`. In this
case, there's no way to know whether the insert was actually successful!

If, however, you want to gracefully handle this in your own application, you
could do something like:

```ruby
document = { _id: Moped::BSON::ObjectId.new, name: "John" }

begin
  session["users"].insert(document)
rescue Moped::Errors::SocketError
  session["users"].find(_id: document[:_id]).upsert(document)
end
```

# Replica Sets

Moped has full support for replica sets including automatic failover and node
discovery.

## Automatic Failover

Moped will automatically retry lost connections and attempt to detect dead
connections before sending an operation. Note, that it will *not* retry
individual operations! For example, these cases will work and not raise any
exceptions:

```ruby
session[:users].insert(name: "John")
# kill primary node and promote secondary
session[:users].insert(name: "John")
session[:users].find.count # => 2.0

# primary node drops our connection
session[:users].insert(name: "John")
```

However, you'll get an operation error in a case like:

```ruby
# primary node goes down while reading the reply
session.with(safe: true)[:users].insert(name: "John")
```

And you'll get a connection error in a case like:

```ruby
# primary node goes down, no new primary available yet
session[:users].insert(name: "John")
```

If your session is running with eventual consistency, read operations will
never raise connection errors as long as any secondary or primary node is
running. The only case where you'll see a connection failure is if a node goes
down while attempting to retrieve more results from a cursor, because cursors
are tied to individual nodes.

When two attempts to connect to a node fail, it will be marked as down. This
removes it from the list of available nodes for `:down_interval` (default 30
seconds). Note that the `:down_interval` only applies to normal operations;
that is, if you ask for a primary node and none is available, all nodes will be
retried. Likewise, if you ask for a secondary node, and no secondary or primary
node is available, all nodes will be retreied.

## Node Discovery

The addresses you pass into your session are used as seeds for setting up
replica set connections. After connection, each seed node will return a list of
other known nodes which will be added to the set.

This information is cached according to the `:refresh_interval` option (default:
5 minutes). That means, e.g., that if you add a new node to your replica set,
it should be represented in Moped within 5 minutes.

require "moped/connection"
require "moped/node"
require "moped/session/context"
require "moped/threaded"
          data = ""
          from_data data
        end

        def from_time(time)
          from_data @@generator.generate(time.to_i)
          !!str.match(/\A\h{24}\Z/i)
        def from_data(data)
          id = allocate
          id.instance_variable_set :@data, data
          id
      def data
        @data ||= @@generator.next
      end

        @@string_format % data.unpack("C12")
        Time.at(data.unpack("N")[0]).utc
          from_data(io.read(12))
        io << data
          @machine_id = Digest::MD5.digest(Socket.gethostname).unpack("N")[0]
        # Return object id data based on the current time, incrementing the
        # object id counter.
          @mutex.lock
          begin
            counter = @counter = (@counter + 1) % 0xFFFFFF
          ensure
            @mutex.unlock rescue nil
          generate(Time.new.to_i, counter)
        # Generate object id data for a given time using the provided +counter+.
        def generate(time, counter = 0)
          [time, @machine_id, Process.pid, counter << 8].pack("N NX lXX NX")
    # @return [Array<String>] the seeds the replica set was initialized with
    # @option options :down_interval number of seconds to wait before attempting
    # to reconnect to a down node. (30)
    #
    # @option options :refresh_interval number of seconds to cache information
    # about a node. (300)
    def initialize(hosts, options)
      @options = {
        down_interval: 30,
        refresh_interval: 300
      }.merge(options)

      @seeds = hosts
      @nodes = hosts.map { |host| Node.new(host) }
    # Refreshes information for each of the nodes provided. The node list
    # defaults to the list of all known nodes.
    #
    # If a node is successfully refreshed, any newly discovered peers will also
    # be refreshed.
    #
    # @return [Array<Node>] the available nodes
    def refresh(nodes_to_refresh = @nodes)
      refreshed_nodes = []
      seen = {}

      # Set up a recursive lambda function for refreshing a node and it's peers.
      refresh_node = ->(node) do
        return if seen[node]
        seen[node] = true

        # Add the node to the global list of known nodes.
        @nodes << node unless @nodes.include?(node)

        begin
          node.refresh

          # This node is good, so add it to the list of nodes to return.
          refreshed_nodes << node unless refreshed_nodes.include?(node)

          # Now refresh any newly discovered peer nodes.
          (node.peers - @nodes).each &refresh_node
        rescue Errors::ConnectionFailure
          # We couldn't connect to the node, so don't do anything with it.
      nodes_to_refresh.each &refresh_node
      refreshed_nodes.to_a
    # Returns the list of available nodes, refreshing 1) any nodes which were
    # down and ready to be checked again and 2) any nodes whose information is
    # out of date.
    #
    # @return [Array<Node>] the list of available nodes.
    def nodes
      # Find the nodes that were down but are ready to be refreshed, or those
      # with stale connection information.
      needs_refresh, available = @nodes.partition do |node|
        (node.down? && node.down_at < (Time.new - @options[:down_interval])) ||
          node.needs_refresh?(Time.new - @options[:refresh_interval])
      end
      # Refresh those nodes.
      available.concat refresh(needs_refresh)
      # Now return all the nodes that are available.
      available.reject &:down?
    end
    # Yields the replica set's primary node to the provided block. This method
    # will retry the block in case of connection errors or replica set
    # reconfiguration.
    #
    # @raises ConnectionFailure when no primary node can be found
    def with_primary(retry_on_failure = true, &block)
      if node = nodes.find(&:primary?)
        begin
          node.ensure_primary do
            return yield node.apply_auth(auth)
        rescue Errors::ConnectionFailure, Errors::ReplicaSetReconfigured
          # Fall through to the code below if our connection was dropped or the
          # node is no longer the primary.
      end
      if retry_on_failure
        # We couldn't find a primary node, so refresh the list and try again.
        refresh
        with_primary(false, &block)
        raise Errors::ConnectionFailure, "Could not connect to a primary node for replica set #{inspect}"
    # Yields a secondary node if available, otherwise the primary node. This
    # method will retry the block in case of connection errors.
    #
    # @raises ConnectionError when no secondary or primary node can be found
    def with_secondary(retry_on_failure = true, &block)
      available_nodes = nodes.shuffle!.partition(&:secondary?).flatten

      while node = available_nodes.shift
        begin
          return yield node.apply_auth(auth)
        rescue Errors::ConnectionFailure
          # That node's no good, so let's try the next one.
          next
      if retry_on_failure
        # We couldn't find a secondary or primary node, so refresh the list and
        # try again.
        refresh
        with_secondary(false, &block)
      else
        raise Errors::ConnectionFailure, "Could not connect to any secondary or primary nodes for replica set #{inspect}"
    private
    def initialize_copy(_)
      @nodes = @nodes.map &:dup
        session.context.insert(database.name, name, documents)
require "timeout"

module Moped
  class Connection

    class TCPSocket < ::TCPSocket
      def self.connect(host, port, timeout)
        Timeout::timeout(timeout) do
          new(host, port).tap do |sock|
            sock.set_encoding 'binary'
          end
        end
      end

      def alive?
        if Kernel::select([self], nil, nil, 0)
          !eof? rescue false
        else
          true
        end
      end

      def write(*args)
        raise Errors::ConnectionFailure, "Socket connection was closed by remote host" unless alive?
        super
      end
    end

    def initialize
      @sock = nil
      @request_id = 0
    end

    def connect(host, port, timeout)
      @sock = TCPSocket.connect host, port, timeout
    end

    def alive?
      connected? ? @sock.alive? : false
    end

    def connected?
      !!@sock
    end

    def disconnect
      @sock.close
    rescue
    ensure
      @sock = nil
    end

    def write(operations)
      buf = ""

      operations.each do |operation|
        operation.request_id = (@request_id += 1)
        operation.serialize(buf)
      end

      @sock.write buf
    end

    def receive_replies(operations)
      operations.map do |operation|
        read if operation.is_a?(Protocol::Query) || operation.is_a?(Protocol::GetMore)
      end
    end

    def read
      reply = Protocol::Reply.allocate

      reply.length,
        reply.request_id,
        reply.response_to,
        reply.op_code,
        reply.flags,
        reply.cursor_id,
        reply.offset,
        reply.count = @sock.read(36).unpack('l5<q<l2<')

      if reply.count == 0
        reply.documents = []
      else
        buffer = StringIO.new(@sock.read(reply.length - 36))

        reply.documents = reply.count.times.map do
          BSON::Document.deserialize(buffer)
        end
      end

      reply
    end
  end
end
      @database    = query_operation.database
      @collection  = query_operation.collection
      @selector    = query_operation.selector
      @cursor_id = 0
      @limit = query_operation.limit
      @limited = @limit > 0

      @options = {
        request_id: query_operation.request_id,
        flags: query_operation.flags,
        limit: query_operation.limit,
        skip: query_operation.skip,
        fields: query_operation.fields
      }
      documents = load
        return kill if limited? && @limit <= 0
        documents = get_more
    def load
      consistency = session.consistency
      @options[:flags] |= [:slave_ok] if consistency == :eventual

      reply, @node = session.context.with_node do |node|
        [node.query(@database, @collection, @selector, @options), node]
      end
      @limit -= reply.count if limited?
      @cursor_id = reply.cursor_id
      @limited
      @cursor_id != 0
    end

    def get_more
      reply = @node.get_more @database, @collection, @cursor_id, @limit

      @limit -= reply.count if limited?
      @cursor_id = reply.cursor_id

      reply.documents
      @node.kill_cursors [@cursor_id]
      session.with(consistency: :strong) do |session|
        session.context.command name, dropDatabase: 1
      end
      session.context.login(name, username, password)
      session.context.logout(name)
      session.context.command name, command
    # Tag applied to unhandled exceptions on a node.
    module SocketError end
    class MongoError < StandardError
    # Exception class for exceptions generated as a direct result of an
    # operation, such as a failed insert or an invalid command.
    class OperationFailure < MongoError; end

    # Exception raised on invalid queries.
    class QueryFailure < MongoError; end

    # Exception raised when authentication fails.
    class AuthenticationFailure < MongoError; end

    # Raised when providing an invalid string from an object id.
    class InvalidObjectId < StandardError
      def initialize(string)
        super("'#{string}' is not a valid object id.")
      end
    end

    # @api private
    #
    # Internal exception raised by Node#ensure_primary and captured by
    # Cluster#with_primary.
    class ReplicaSetReconfigured < StandardError; end
module Moped
  class Node

    attr_reader :address
    attr_reader :resolved_address
    attr_reader :ip_address
    attr_reader :port

    attr_reader :peers
    attr_reader :timeout

    def initialize(address)
      @address = address

      host, port = address.split(":")
      @ip_address = ::Socket.getaddrinfo(host, nil, ::Socket::AF_INET, ::Socket::SOCK_STREAM).first[3]
      @port = port.to_i
      @resolved_address = "#{@ip_address}:#{@port}"

      @timeout = 5
    end

    def command(database, cmd, options = {})
      operation = Protocol::Command.new(database, cmd, options)

      process(operation) do |reply|
        result = reply.documents[0]

        raise Errors::OperationFailure.new(
          operation, result
        ) if result["ok"] != 1 || result["err"] || result["errmsg"]

        result
      end
    end

    def kill_cursors(cursor_ids)
      process Protocol::KillCursors.new(cursor_ids)
    end

    def get_more(database, collection, cursor_id, limit)
      process Protocol::GetMore.new(database, collection, cursor_id, limit)
    end

    def remove(database, collection, selector, options = {})
      process Protocol::Delete.new(database, collection, selector, options)
    end

    def update(database, collection, selector, change, options = {})
      process Protocol::Update.new(database, collection, selector, change, options)
    end

    def insert(database, collection, documents)
      process Protocol::Insert.new(database, collection, documents)
    end

    def query(database, collection, selector, options = {})
      operation = Protocol::Query.new(database, collection, selector, options)

      process operation do |reply|
        if reply.flags.include? :query_failure
          raise Errors::QueryFailure.new(operation, reply.documents.first)
        end

        reply
      end
    end

    # @return [true/false] whether the node needs to be refreshed.
    def needs_refresh?(time)
      !@refreshed_at || @refreshed_at < time
    end

    def primary?
      @primary
    end

    def secondary?
      @secondary
    end

    # Refresh information about the node, such as it's status in the replica
    # set and it's known peers.
    #
    # Returns nothing.
    # Raises Errors::ConnectionFailure if the node cannot be reached
    # Raises Errors::ReplicaSetReconfigured if the node is no longer a primary node and
    #   refresh was called within an +#ensure_primary+ block.
    def refresh
      info = command "admin", ismaster: 1

      @refreshed_at = Time.now
      primary = true if info["ismaster"]
      secondary = true if info["secondary"]

      peers = []
      peers.push info["primary"] if info["primary"]
      peers.concat info["hosts"] if info["hosts"]
      peers.concat info["passives"] if info["passives"]
      peers.concat info["arbiters"] if info["arbiters"]

      @peers = peers.map { |peer| Node.new(peer) }
      @primary, @secondary = primary, secondary

      if !primary && Threaded.executing?(:ensure_primary)
        raise Errors::ReplicaSetReconfigured, "#{inspect} is no longer the primary node."
      end
    end

    attr_reader :down_at

    def down?
      @down_at
    end

    # Set a flag on the node for the duration of provided block so that an
    # exception is raised if the node is no longer the primary node.
    #
    # Returns nothing.
    def ensure_primary
      Threaded.begin :ensure_primary
      yield
    ensure
      Threaded.end :ensure_primary
    end

    # Yields the block if a connection can be established, retrying when a
    # connection error is raised.
    #
    # @raises ConnectionFailure when a connection cannot be established.
    def ensure_connected
      # Don't run the reconnection login if we're already inside an
      # +ensure_connected+ block.
      return yield if Threaded.executing? :connection
      Threaded.begin :connection

      retry_on_failure = true

      begin
        connect unless connected?
        yield
      rescue Errors::ReplicaSetReconfigured
        # Someone else wrapped this in an #ensure_primary block, so let the
        # reconfiguration exception bubble up.
        raise
      rescue Errors::OperationFailure, Errors::AuthenticationFailure, Errors::QueryFailure
        # These exceptions are "expected" in the normal course of events, and
        # don't necessitate disconnecting.
        raise
      rescue Errors::ConnectionFailure
        disconnect

        if retry_on_failure
          # Maybe there was a hiccup -- try reconnecting one more time
          retry_on_failure = false
          retry
        else
          # Nope, we failed to connect twice. Flag the node as down and re-raise
          # the exception.
          down!
          raise
        end
      rescue
        # Looks like we got an unexpected error, so we'll clean up the connection
        # and re-raise the exception.
        disconnect
        raise $!.extend(Errors::SocketError)
      end
    ensure
      Threaded.end :connection
    end

    def pipeline
      Threaded.begin :pipeline

      begin
        yield
      ensure
        Threaded.end :pipeline
      end

      flush unless Threaded.executing? :pipeline
    end

    def apply_auth(credentials)
      unless auth == credentials
        logouts = auth.keys - credentials.keys

        logouts.each do |database|
          logout database
        end

        credentials.each do |database, (username, password)|
          login(database, username, password) unless auth[database] == [username, password]
        end
      end

      self
    end

    def ==(other)
      resolved_address == other.resolved_address
    end
    alias eql? ==

    def hash
      [ip_address, port].hash
    end

    private

    def auth
      @auth ||= {}
    end

    def login(database, username, password)
      getnonce = Protocol::Command.new(database, getnonce: 1)
      connection.write [getnonce]
      result = connection.read.documents.first
      raise Errors::OperationFailure.new(getnonce, result) unless result["ok"] == 1

      authenticate = Protocol::Commands::Authenticate.new(database, username, password, result["nonce"])
      connection.write [authenticate]
      result = connection.read.documents.first
      raise Errors::AuthenticationFailure.new(authenticate, result) unless result["ok"] == 1

      auth[database] = [username, password]
    end

    def logout(database)
      command = Protocol::Command.new(database, logout: 1)
      connection.write [command]
      result = connection.read.documents.first
      raise Errors::OperationFailure.new(command, result) unless result["ok"] == 1

      auth.delete(database)
    end

    def initialize_copy(_)
      @connection = nil
    end

    def connection
      @connection ||= Connection.new
    end

    def disconnect
      auth.clear
      connection.disconnect
    end

    def connected?
      connection.connected?
    end

    # Mark the node as down.
    #
    # Returns nothing.
    def down!
      @down_at = Time.new

      disconnect
    end

    # Connect to the node.
    #
    # Returns nothing.
    # Raises Moped::ConnectionError if the connection times out.
    # Raises Moped::ConnectionError if the server is unavailable.
    def connect
      connection.connect ip_address, port, timeout
      @down_at = nil

      refresh
    rescue Timeout::Error
      raise Errors::ConnectionFailure, "Timed out connection to Mongo on #{address}"
    rescue Errno::ECONNREFUSED
      raise Errors::ConnectionFailure, "Could not connect to Mongo on #{address}"
    end

    def process(operation, &callback)
      if Threaded.executing? :pipeline
        queue.push [operation, callback]
      else
        flush([[operation, callback]])
      end
    end

    def queue
      Threaded.stack(:pipelined_operations)
    end

    def flush(ops = queue)
      operations, callbacks = ops.transpose

      logging(operations) do
        ensure_connected do
          connection.write operations
          replies = connection.receive_replies(operations)

          replies.zip(callbacks).map do |reply, callback|
            callback ? callback[reply] : reply
          end.last
        end
      end
    ensure
      ops.clear
    end

    def logging(operations)
      instrument_start = (logger = Moped.logger) && logger.debug? && Time.new
      yield
    ensure
      log_operations(logger, operations, Time.new - instrument_start) if instrument_start && !$!
    end

    def log_operations(logger, ops, duration)
      prefix  = "  MOPED: #{address} "
      indent  = " "*prefix.length
      runtime = (" (%.1fms)" % duration)

      if ops.length == 1
        logger.debug prefix + ops.first.log_inspect + runtime
      else
        first, *middle, last = ops

        logger.debug prefix + first.log_inspect
        middle.each { |m| logger.debug indent + m.log_inspect }
        logger.debug indent + last.log_inspect + runtime
      end
    end

  end
end
      # @param [Hash] additional query options
      def initialize(database, command, options = {})
        super database, :$cmd, command, options.merge(limit: -1)
      limit(-1).each.first
      cursor = Cursor.new(session, operation)

        session.context.update operation.database,
          operation.collection,
          operation.selector,
          change,
          flags: flags
        session.context.remove operation.database,
          operation.collection,
          operation.selector,
          flags: [:remove_first]
        session.context.remove operation.database,
          operation.collection,
          operation.selector
    # @private
    # @return [Context] this session's context
    attr_reader :context

      @cluster = Cluster.new(seeds, {})
      @context = Context.new(self)
    # @return [:strong, :eventual] the session's consistency
    def consistency
      options[:consistency]
    end

    #   session[:people].find.one # => { :name => "John" }
      session.instance_variable_set(:@cluster, cluster.dup)
    private

    def current_database
      return @current_database if defined? @current_database

      if database = options[:database]
        set_current_database(database)
        raise "No database set for session. Call #use or #with before accessing the database"
    def initialize_copy(_)
      @context = Context.new(self)
      @options = @options.dup
        remove_instance_variable :@current_database
module Moped
  class Session

    # @api private
    class Context
      extend Forwardable

      def initialize(session)
        @session = session
      end

      delegate :safety => :@session
      delegate :safe?  => :@session
      delegate :consistency  => :@session
      delegate :cluster  => :@session

      def login(database, username, password)
        cluster.auth[database.to_s] = [username, password]
      end

      def logout(database)
        cluster.auth.delete database.to_s
      end

      def query(database, collection, selector, options = {})
        if consistency == :eventual
          options[:flags] ||= []
          options[:flags] |= [:slave_ok]
        end

        with_node do |node|
          node.query(database, collection, selector, options)
        end
      end

      def command(database, command)
        options = consistency == :eventual ? { :flags => [:slave_ok] } : {}
        with_node do |node|
          node.command(database, command, options)
        end
      end

      def insert(database, collection, documents)
        with_node do |node|
          if safe?
            node.pipeline do
              node.insert(database, collection, documents)
              node.command("admin", { getlasterror: 1 }.merge(safety))
            end
          else
            node.insert(database, collection, documents)
          end
        end
      end

      def update(database, collection, selector, change, options = {})
        with_node do |node|
          if safe?
            node.pipeline do
              node.update(database, collection, selector, change, options)
              node.command("admin", { getlasterror: 1 }.merge(safety))
            end
          else
            node.update(database, collection, selector, change, options)
          end
        end
      end

      def remove(database, collection, selector, options = {})
        with_node do |node|
          if safe?
            node.pipeline do
              node.remove(database, collection, selector, options)
              node.command("admin", { getlasterror: 1 }.merge(safety))
            end
          else
            node.remove(database, collection, selector, options)
          end
        end
      end

      def get_more(*args)
        raise NotImplementedError, "#get_more cannot be called on Context; it must be called directly on a node"
      end

      def kill_cursors(*args)
        raise NotImplementedError, "#kill_cursors cannot be called on Context; it must be called directly on a node"
      end

      def with_node
        if consistency == :eventual
          cluster.with_secondary do |node|
            yield node
          end
        else
          cluster.with_primary do |node|
            yield node
          end
        end
      end

    end

  end
end
module Moped

  # This module contains logic for easy access to objects that have a lifecycle
  # on the current thread.
  #
  # Extracted from Mongoid's +Threaded+ module.
  #
  # @api private
  module Threaded
    extend self

    # Begin a thread-local stack for +name+.
    def begin(name)
      stack(name).push true
    end

    # @return [Boolean] whether the stack is being executed
    def executing?(name)
      !stack(name).empty?
    end

    # End the thread-local stack for +name+.
    def end(name)
      stack(name).pop
    end

    # @return [Array] a named, thread-local stack.
    def stack(name)
      Thread.current["[moped]:#{name}-stack"] ||= []
    end
  end
end
      session.new do |session|
        1000.times do
          session[:people].insert(name: "John")
          session[:people].find.one
        end
      session.new do |session|
        1000.times do
          session[:people].insert(name: "John")
          session[:people].find.one
        end
    [78, 77, 102, 52, 59, 57, 182, 132, 7, 0, 0, 1].pack("C12")
      it "initializes with the string's bytes" do
        Moped::BSON::ObjectId.should_receive(:from_data).with(bytes)
  describe "#from_time" do
    it "sets the generation time" do
      time = Time.at((Time.now.utc - 64800).to_i).utc
      Moped::BSON::ObjectId.from_time(time).generation_time.should == time
  end
  describe "#initialize" do
        Moped::BSON::ObjectId.from_data(bytes).should == Moped::BSON::ObjectId.from_data(bytes)
        Moped::BSON::ObjectId.from_data(bytes).should eql Moped::BSON::ObjectId.from_data(bytes)
        Moped::BSON::ObjectId.from_data(bytes).hash.should eq Moped::BSON::ObjectId.from_data(bytes).hash
      Moped::BSON::ObjectId.from_data(bytes).to_s.should eq "4e4d66343b39b68407000001"
describe Moped::Cluster, replica_set: true do
  let(:replica_set) do
    Moped::Cluster.new(seeds, {})
  context "when the replica set hasn't connected yet" do
    describe "#with_primary" do
      it "connects and yields the primary node" do
        replica_set.with_primary do |node|
          node.address.should eq @primary.address
        end
      end
    describe "#with_secondary" do
      it "connects and yields a secondary node" do
        replica_set.with_secondary do |node|
          @secondaries.map(&:address).should include node.address
        end
    context "and the primary is down" do
        @primary.stop
      describe "#with_primary" do
        it "raises a connection error" do
          lambda do
            replica_set.with_primary do |node|
              node.command "admin", ping: 1
            end
          end.should raise_exception(Moped::Errors::ConnectionFailure)
      describe "#with_secondary" do
        it "connects and yields a secondary node" do
          replica_set.with_secondary do |node|
            @secondaries.map(&:address).should include node.address
          end
    end
    context "and a single secondary is down" do
      before do
        @secondaries.first.stop
      end
      describe "#with_primary" do
        it "connects and yields the primary node" do
          replica_set.with_primary do |node|
            node.address.should eq @primary.address
          end
      end
      describe "#with_secondary" do
        it "connects and yields a secondary node" do
          replica_set.with_secondary do |node|
            node.address.should eq @secondaries.last.address
          end
    end
    context "and all secondaries are down" do
      before do
        @secondaries.each &:stop
      end
      describe "#with_primary" do
        it "connects and yields the primary node" do
          replica_set.with_primary do |node|
            node.address.should eq @primary.address
          end
      end
      describe "#with_secondary" do
        it "connects and yields the primary node" do
          replica_set.with_secondary do |node|
            node.address.should eq @primary.address
          end
  context "when the replica set is connected" do
    before do
      replica_set.refresh
    describe "#with_primary" do
      it "connects and yields the primary node" do
        replica_set.with_primary do |node|
          node.address.should eq @primary.address
        end
    describe "#with_secondary" do
      it "connects and yields a secondary node" do
        replica_set.with_secondary do |node|
          @secondaries.map(&:address).should include node.address
        end
    context "and the primary is down" do
      before do
        @primary.stop
      end

      describe "#with_primary" do
        it "raises a connection error" do
          lambda do
            replica_set.with_primary do |node|
              node.command "admin", ping: 1
            end
          end.should raise_exception(Moped::Errors::ConnectionFailure)
      describe "#with_secondary" do
        it "connects and yields a secondary node" do
          replica_set.with_secondary do |node|
            @secondaries.map(&:address).should include node.address
          end
    end
    context "and a single secondary is down" do
        @secondaries.first.stop
      describe "#with_primary" do
        it "connects and yields the primary node" do
          replica_set.with_primary do |node|
            node.address.should eq @primary.address
          end
        end
      describe "#with_secondary" do
        it "connects and yields a secondary node" do
          replica_set.with_secondary do |node|
            node.command "admin", ping: 1
            node.address.should eq @secondaries.last.address
          end
        end
    context "and all secondaries are down" do
        @secondaries.each &:stop
      describe "#with_primary" do
        it "connects and yields the primary node" do
          replica_set.with_primary do |node|
            node.address.should eq @primary.address
      end
      describe "#with_secondary" do
        it "connects and yields the primary node" do
          replica_set.with_secondary do |node|
            node.command "admin", ping: 1
            node.address.should eq @primary.address
          end
      end
    end
  end
  context "with down interval" do
    let(:replica_set) do
      Moped::Cluster.new(seeds, { down_interval: 5 })
    end

    context "and all secondaries are down" do
      before do
        replica_set.refresh
        @secondaries.each &:stop
        replica_set.refresh
      describe "#with_secondary" do
        it "connects and yields the primary node" do
          replica_set.with_secondary do |node|
            node.command "admin", ping: 1
            node.address.should eq @primary.address
          end
      end
      context "when a secondary node comes back up" do
        before do
          @secondaries.each &:restart
        describe "#with_secondary" do
          it "connects and yields the primary node" do
            replica_set.with_secondary do |node|
              node.command "admin", ping: 1
              node.address.should eq @primary.address
            end
          end
        context "and the node is ready to be retried" do
          it "connects and yields the secondary node" do
            Time.stub(:new).and_return(Time.now + 10)
            replica_set.with_secondary do |node|
              node.command "admin", ping: 1
              @secondaries.map(&:address).should include node.address
            end
          end
  end
  context "with only primary provided as a seed" do
    let(:replica_set) do
      Moped::Cluster.new([@primary.address], {})
    end
    describe "#with_primary" do
      it "connects and yields the primary node" do
        replica_set.with_primary do |node|
          node.address.should eq @primary.address
      end
    end
    describe "#with_secondary" do
      it "connects and yields a secondary node" do
        replica_set.with_secondary do |node|
          @secondaries.map(&:address).should include node.address
    end
  end
  context "with only primary provided as a seed" do
    let(:replica_set) do
      Moped::Cluster.new([@secondaries[0].address], {})
    end
    describe "#with_primary" do
      it "connects and yields the primary node" do
        replica_set.with_primary do |node|
          node.address.should eq @primary.address
      end
    end
    describe "#with_secondary" do
      it "connects and yields a secondary node" do
        replica_set.with_secondary do |node|
          @secondaries.map(&:address).should include node.address
end
describe Moped::Cluster, "authentication", mongohq: :auth do
  let(:session) do
    Support::MongoHQ.auth_session(false)
  end
  describe "logging in with valid credentials" do
    it "logs in and processes commands" do
      session.login *Support::MongoHQ.auth_credentials
      session.command(ping: 1).should eq("ok" => 1)
  describe "logging in with invalid credentials" do
    it "raises an AuthenticationFailure exception" do
      session.login "invalid-user", "invalid-password"

      lambda do
        session.command(ping: 1)
      end.should raise_exception(Moped::Errors::AuthenticationFailure)
  end
  describe "logging in with valid credentials and then logging out" do
      session.login *Support::MongoHQ.auth_credentials
      session.command(ping: 1).should eq("ok" => 1)
    it "logs out" do
      lambda do
        session.command dbStats: 1
      end.should_not raise_exception

      session.logout

      lambda do
        session.command dbStats: 1
      end.should raise_exception(Moped::Errors::OperationFailure)
    Moped::Session.new %w[127.0.0.1:27017], database: "moped_test"
  let(:scope) { object_id }
    before do
      session.drop
      session.command create: "users"
    it "drops the collection" do
      result = session[:users].drop
      result["ns"].should eq "moped_test.users"
    it "inserts a single document" do
      document = { "_id" => Moped::BSON::ObjectId.new, "scope" => scope }
      session[:users].insert(document)
      session[:users].find(document).one.should eq document
    it "insert multiple documents" do
      documents = [
        { "_id" => Moped::BSON::ObjectId.new, "scope" => scope },
        { "_id" => Moped::BSON::ObjectId.new, "scope" => scope }
      ]
      session[:users].insert(documents)
      session[:users].find(scope: scope).entries.should eq documents
  let(:session) do
    Moped::Session.new %w[127.0.0.1:27017], database: "moped_test"
  let(:indexes) do
    session[:users].indexes
  before do
    indexes.drop
    context "when called without extra options" do
      it "creates an index with no options" do
        indexes.create name: 1
        indexes[name: 1].should_not be_nil
    context "when called with extra options" do
      it "creates an index with the extra options" do
        indexes.create({name: 1}, {unique: true, dropDups: true})
        index = indexes[name: 1]
        index["unique"].should be_true
        index["dropDups"].should be_true
    context "when provided a key" do
      it "drops the index" do
        indexes.create name: 1
        indexes.drop(name: 1).should be_true
    context "when not provided a key" do
      it "drops all indexes" do
        indexes.create name: 1
        indexes.create age: 1
        indexes.drop
        indexes[age: 1].should be_nil
require "spec_helper"

describe Moped::Node, replica_set: true do
  let(:replica_set_node) do
    @replica_set.nodes.first
  end

  let(:node) do
    Moped::Node.new(replica_set_node.address)
  end

  describe "#ensure_connected" do
    context "when node is running" do
      it "processes the block" do
        node.ensure_connected do
          node.command("admin", ping: 1)
        end.should eq("ok" => 1)
      end
    end

    context "when node is not running" do
      before do
        replica_set_node.stop
      end

      it "raises a connection error" do
        lambda do
          node.ensure_connected do
            node.command("admin", ping: 1)
          end
        end.should raise_exception(Moped::Errors::ConnectionFailure)
      end

      it "marks the node as down" do
        node.ensure_connected {} rescue nil
        node.should be_down
      end
    end

    context "when node is connected but connection is dropped" do
      before do
        node.ensure_connected do
          node.command("admin", ping: 1)
        end

        replica_set_node.hiccup
      end

      it "reconnects without raising an exception" do
        node.ensure_connected do
          node.command("admin", ping: 1)
        end.should eq("ok" => 1)
      end
    end
  end

end
  shared_examples_for "Query" do
    let(:scope) do
      object_id
    end
    before do
      users.find.remove_all
    let(:documents) do
      [
        { "_id" => Moped::BSON::ObjectId.new, "scope" => scope },
        { "_id" => Moped::BSON::ObjectId.new, "scope" => scope }
      ]
    it "raises a query failure exception for invalid queries" do
      lambda do
        users.find("age" => { "$in" => nil }).first
      end.should raise_exception(Moped::Errors::QueryFailure)
    end
    describe "#limit" do
      it "limits the query" do
        users.insert(documents)
        users.find(scope: scope).limit(1).to_a.should eq [documents.first]
      end
    describe "#skip" do
      it "skips +n+ documents" do
        users.insert(documents)
        users.find(scope: scope).skip(1).to_a.should eq [documents.last]
      end
    describe "#sort" do
      let(:documents) do
        [
          { "_id" => Moped::BSON::ObjectId.new, "scope" => scope, "n" => 0 },
          { "_id" => Moped::BSON::ObjectId.new, "scope" => scope, "n" => 1 }
        ]
      end
      it "sorts the results" do
        users.insert(documents)
        users.find(scope: scope).sort(n: -1).to_a.should eq documents.reverse
      end
    describe "#distinct" do
      let(:documents) do
        [
          { count: 0, scope: scope },
          { count: 1, scope: scope },
          { count: 1, scope: scope }
        ]
      end

      it "returns distinct values for +key+" do
        users.insert(documents)
        users.find(scope: scope).distinct(:count).should =~ [0, 1]
      end
    describe "#select" do
      let(:documents) do
        [
          { "scope" => scope, "n" => 0 },
          { "scope" => scope, "n" => 1 }
        ]
      end
      it "changes the fields returned" do
        users.insert(documents)
        users.find(scope: scope).select(_id: 0).to_a.should eq documents
      end
    describe "#one" do
      before do
        users.insert(documents)
      end
      it "returns the first matching document" do
        users.find(scope: scope).one.should eq documents.first
      end
      it "respects #skip" do
        users.find(scope: scope).skip(1).one.should eq documents.last
      end
      it "respects #sort" do
        users.find(scope: scope).sort(_id: -1).one.should eq documents.last
    describe "#explain" do
      context "when a sort exists" do
        it "updates to a mongo advanced selector" do
          stats = Support::Stats.collect do
            users.find(scope: scope).sort(_id: 1).explain
          end

          operation = stats[node_for_reads].grep(Moped::Protocol::Query).last
          operation.selector.should eq(
            "$query" => { scope: scope },
            "$explain" => true,
            "$orderby" => { _id: 1 }
          )
        end
      end
      context "when no sort exists" do
        it "updates to a mongo advanced selector" do
          stats = Support::Stats.collect do
            users.find(scope: scope).explain
          end

          operation = stats[node_for_reads].grep(Moped::Protocol::Query).last
          operation.selector.should eq(
            "$query" => { scope: scope },
            "$explain" => true,
            "$orderby" => {}
          )
        end
    describe "#each" do
      it "yields each document" do
        users.insert(documents)
        users.find(scope: scope).each.with_index do |document, index|
          document.should eq documents[index]
        end
      end
      context "with a limit" do
        it "closes open cursors" do
          users.insert(100.times.map { Hash["scope" => scope] })
          stats = Support::Stats.collect do
            users.find(scope: scope).limit(5).entries
          end
          stats[node_for_reads].grep(Moped::Protocol::KillCursors).count.should eq 1
        end
      context "without a limit" do
        it "fetches more" do
          users.insert(102.times.map { Hash["scope" => scope] })
          stats = Support::Stats.collect do
            users.find(scope: scope).entries
          end
          stats[node_for_reads].grep(Moped::Protocol::GetMore).count.should eq 1
        end
    describe "#count" do
      let(:documents) do
        [
          { "_id" => Moped::BSON::ObjectId.new, "scope" => scope },
          { "_id" => Moped::BSON::ObjectId.new, "scope" => scope },
          { "_id" => Moped::BSON::ObjectId.new }
        ]
      end
      it "returns the number of matching document" do
        users.insert(documents)
        users.find(scope: scope).count.should eq 2
      end
    describe "#update" do
      it "updates the first matching document" do
        users.insert(documents)
        users.find(scope: scope).update("$set" => { "updated" => true })
        users.find(scope: scope, updated: true).count.should eq 1
      end
    end
    describe "#update_all" do
      it "updates all matching documents" do
        users.insert(documents)
        users.find(scope: scope).update_all("$set" => { "updated" => true })
        users.find(scope: scope, updated: true).count.should eq 2
      end
    describe "#upsert" do
      context "when a document exists" do
        before do
          users.insert(scope: scope, counter: 1)
        end
        it "updates the document" do
          users.find(scope: scope).upsert("$inc" => { counter: 1 })
          users.find(scope: scope).one["counter"].should eq 2
        end
      end
      context "when no document exists" do
        it "inserts a document" do
          users.find(scope: scope).upsert("$inc" => { counter: 1 })
          users.find(scope: scope).one["counter"].should eq 1
        end
      end
    describe "#remove" do
      it "removes the first matching document" do
        users.insert(documents)
        users.find(scope: scope).remove
        users.find(scope: scope).count.should eq 1
      end
    end
    describe "#remove_all" do
      it "removes all matching documents" do
        users.insert(documents)
        users.find(scope: scope).remove_all
        users.find(scope: scope).count.should eq 0
      end
  context "with a local connection" do
    let(:session) do
      Moped::Session.new %w[127.0.0.1:27017], database: "moped_test"
    let(:users) { session[:users] }
    let(:node_for_reads) { :primary }
    include_examples "Query"
    describe "#each" do
      context "with a limit and large result set" do
        it "gets more and closes cursors" do
          11.times do
            users.insert(scope: scope, large_field: "a"*1_000_000)
          end
          stats = Support::Stats.collect do
            users.find(scope: scope).limit(10).entries
          end
          stats[:primary].grep(Moped::Protocol::GetMore).count.should eq 1
          stats[:primary].grep(Moped::Protocol::KillCursors).count.should eq 1
        end
      end
  context "with a remote connection", mongohq: :auth do
    before :all do
      @session = Support::MongoHQ.auth_session
    let(:users) { @session[:users] }
    let(:node_for_reads) { :primary }

    include_examples "Query"
  context "with a remote replica set connection with eventual consistency", mongohq: :replica_set do
    before :all do
      @session = Support::MongoHQ.replica_set_session.with(safe: true, consistency: :eventual)
      @session.command ping: 1
    end
    let(:users) { @session[:users] }
    let(:node_for_reads) { :secondary }
    include_examples "Query"
  context "with a remote replica set connection with strong consistency", mongohq: :replica_set do
    before :all do
      @session = Support::MongoHQ.replica_set_session.with(safe: true, consistency: :strong)
    end
    let(:users) { @session[:users] }
    let(:node_for_reads) { :primary }
    include_examples "Query"
  end
  context "with a local replica set w/ failover", replica_set: true do
    let(:session) do
      Moped::Session.new seeds, database: "moped_test"
    let(:scope) do
      object_id
    end
      # Force connection before recording stats
      session.command ping: 1
    context "and running with eventual consistency" do
      it "queries a secondary node" do
        stats = Support::Stats.collect do
          session.with(consistency: :eventual)[:users].find(scope: scope).entries
        end
        stats[:secondary].grep(Moped::Protocol::Query).count.should eq 1
        stats[:primary].should be_empty
      end

      it "sets the slave ok flag" do
        stats = Support::Stats.collect do
          session.with(consistency: :eventual)[:users].find(scope: scope).one
        end

        query = stats[:secondary].grep(Moped::Protocol::Query).first
        query.flags.should include :slave_ok
      end
      context "and no secondaries are available" do
        before do
          @secondaries.each &:stop
        end
        it "queries the primary node" do
          stats = Support::Stats.collect do
            session.with(consistency: :eventual)[:users].find(scope: scope).entries
          end
          stats[:primary].grep(Moped::Protocol::Query).count.should eq 1
        end
      end
    context "and running with strong consistency" do
      it "queries the primary node" do
        stats = Support::Stats.collect do
          session.with(consistency: :strong)[:users].find(scope: scope).entries
        end

        stats[:primary].grep(Moped::Protocol::Query).count.should eq 1
        stats[:secondary].should be_empty
      end
      it "does not set the slave ok flag" do
        stats = Support::Stats.collect do
          session.with(consistency: :strong)[:users].find(scope: scope).one
        end

        query = stats[:primary].grep(Moped::Protocol::Query).first
        query.flags.should_not include :slave_ok
      end
    Moped::Session.new %w[127.0.0.1:27017], database: "moped_test"
    it "changes the current database" do
      session.use "moped_test_2"
      session.command(dbStats: 1)["db"].should eq "moped_test_2"
      it "returns the value from the block" do
        session.with { :value }.should eq :value
      it "yields a session with the provided options" do
        session.with(safe: true) do |safe|
          safe.options[:safe].should eq true
      it "does not modify the original session" do
        session.with(database: "other") do |safe|
          session.options[:database].should eq "moped_test"
      it "returns a session with the provided options" do
        safe = session.with(safe: true)
        safe.options[:safe].should eq true
      it "does not modify the original session" do
        other = session.with(database: "other")
        session.options[:database].should eq "moped_test"
    it "drops the current database" do
      session.with(database: "moped_test_2") do |session|
        session.drop.should eq("dropped" => "moped_test_2", "ok" => 1)
  describe "#command" do
    it "runs the command on the current database" do
      session.with(database: "moped_test_2") do |session|
        session.command(dbStats: 1)["db"].should eq "moped_test_2"
require "rspec"
require "support/mongohq"
require "support/replica_set_simulator"
require "support/stats"
RSpec.configure do |config|
  Support::Stats.install!

  config.include Support::ReplicaSetSimulator::Helpers, replica_set: true

  config.filter_run_excluding mongohq: ->(value) do
    return true if value == :replica_set && !Support::MongoHQ.replica_set_configured?
    return true if value == :auth && !Support::MongoHQ.auth_node_configured?
  end

  unless Support::MongoHQ.replica_set_configured? || Support::MongoHQ.auth_node_configured?
    $stderr.puts Support::MongoHQ.message
  end
end
module Support
  module MongoHQ
    extend self

    def replica_set_configured?
      ENV["MONGOHQ_REPL_PASS"]
    end

    def replica_set_seeds
      [ENV["MONGOHQ_REPL_1_URL"], ENV["MONGOHQ_REPL_2_URL"]]
    end

    def replica_set_credentials
      [ENV["MONGOHQ_REPL_USER"], ENV["MONGOHQ_REPL_PASS"]]
    end

    def replica_set_database
      ENV["MONGOHQ_REPL_NAME"]
    end

    def replica_set_session(auth = true)
      session = Moped::Session.new replica_set_seeds, database: replica_set_database
      session.login *replica_set_credentials if auth
      session
    end

    def auth_seeds
      [ENV["MONGOHQ_SINGLE_URL"]]
    end

    def auth_node_configured?
      ENV["MONGOHQ_SINGLE_PASS"]
    end

    def auth_credentials
      [ENV["MONGOHQ_SINGLE_USER"], ENV["MONGOHQ_SINGLE_PASS"]]
    end

    def auth_database
      ENV["MONGOHQ_SINGLE_NAME"]
    end

    def auth_session(auth = true)
      session = Moped::Session.new auth_seeds, database: auth_database
      session.login *auth_credentials if auth
      session
    end

    def message
      %Q{
      ---------------------------------------------------------------------
      Moped runs specs for authentication and replica sets against MongoHQ.

      If you want to run these specs and need the credentials, contact
      durran at gmail dot com.
      ---------------------------------------------------------------------
      }
    end

  end
end
module Support

  # This is a helper class for testing replica sets. It works by starting up a
  # TCP server socket for each desired node. It then proxies all traffic
  # between a real mongo instance and the client app, with the exception of
  # ismaster commands, which it returns simulated responses for.
  class ReplicaSetSimulator

    module Helpers
      def self.included(context)
        context.before :all do
          @replica_set = ReplicaSetSimulator.new
          @replica_set.start

          @primary, @secondaries = @replica_set.initiate
        end

        context.after :all do
          @replica_set.stop
        end

        context.after :each do
          @replica_set.nodes.each &:restart
        end

        context.let :seeds do
          @replica_set.nodes.map &:address
        end
      end
    end

    attr_reader :nodes
    attr_reader :manager

    def initialize(nodes = 3)
      @nodes = nodes.times.map { Node.new(self) }
      @manager = ConnectionManager.new(@nodes)
      @mongo = TCPSocket.new "127.0.0.1", 27017
    end

    # Start the mock replica set.
    def start
      @nodes.each &:start
      @worker = Thread.start do
        Thread.abort_on_exception = true
        catch(:shutdown) do
          loop do
            Moped.logger.debug "replica_set: waiting for next client"
            server, client = @manager.next_client

            if server
              Moped.logger.debug "replica_set: proxying incoming request to mongo"
              server.proxy(client, @mongo)
            else
              Moped.logger.debug "replica_set: no requests; passing"
              Thread.pass
            end
          end
        end
      end
    end

    # Pick a node to be master, and mark the rest as secondary
    def initiate
      primary, *secondaries = @nodes.shuffle

      primary.promote
      secondaries.each &:demote

      return primary, secondaries
    end

    # Shut down the mock replica set.
    def stop
      @manager.shutdown
      @nodes.each &:stop
    end

    class Node

      attr_reader :port
      attr_reader :host

      def initialize(set)
        @set = set
        @primary = false
        @secondary = false

        server = TCPServer.new 0
        @host = Socket.gethostname
        @port = server.addr[1]
        server.close
      end

      def ==(other)
        @host == other.host && @port == other.port
      end

      def address
        "#{@host}:#{@port}"
      end

      def primary?
        @primary
      end

      def secondary?
        @secondary
      end

      def status
        {
          "ismaster" => @primary,
          "secondary" => @secondary,
          "hosts" => @set.nodes.map(&:address),
          "me" => address,
          "maxBsonObjectSize" => 16777216,
          "ok" => 1.0
        }
      end

      def status_reply
        reply = Moped::Protocol::Reply.new
        reply.count = 1
        reply.documents = [status]
        reply
      end

      OP_QUERY = 2004
      OP_GETMORE = 2005

      # Stop and start the node.
      def restart
        stop
        start
      end

      # Start the node.
      def start
        @server = TCPServer.new @port
      end

      # Stop the node.
      def stop
        if @server
          hiccup

          @server.close
          @server = nil
        end
      end
      alias close stop

      def accept
        to_io.accept
      end

      def closed?
        !@server || @server.closed?
      end

      def to_io
        @server
      end

      # Mark this node as secondary.
      def demote
        @primary = false
        @secondary = true

        hiccup
      end

      def hiccup
        @set.manager.close_clients_for(self)
      end

      # Mark this node as primary. This also closes any open connections.
      def promote
        @primary = true
        @secondary = false

        hiccup
      end

      # Proxies a single message from client to the mongo connection.
      def proxy(client, mongo)
        incoming_message = client.read(16)
        length, op_code = incoming_message.unpack("l<x8l<")
        incoming_message << client.read(length - 16)

        if op_code == OP_QUERY && ismaster_command?(incoming_message)
          # Intercept the ismaster command and send our own reply.
          client.write status_reply
        else
          # This is a normal command, so proxy it to the real mongo instance.
          mongo.write incoming_message

          if op_code == OP_QUERY || op_code == OP_GETMORE
            outgoing_message = mongo.read(4)
            length, = outgoing_message.unpack('l<')
            outgoing_message << mongo.read(length - 4)

            client.write outgoing_message
          end
        end
      end

      private

      # Checks a message to see if it's an `ismaster` query.
      def ismaster_command?(incoming_message)
        data = StringIO.new(incoming_message)
        data.read(20) # header and flags
        data.gets("\x00") # collection name
        data.read(8) # skip/limit

        selector = Moped::BSON::Document.deserialize(data)
        selector == { "ismaster" => 1 }
      end
    end

    class ConnectionManager

      def initialize(servers)
        @timeout = 0.1
        @servers = servers
        @clients = []
      end

      def shutdown
        @servers.each &:close
        @clients.each &:close
        @shutdown = true
      end

      def next_client
        throw :shutdown if @shutdown

        begin
          servers = @servers.reject &:closed?
          clients =  @clients.reject &:closed?
          Moped.logger.debug "replica_set: selecting on connections"
          readable, _, errors = Kernel.select(servers + clients, nil, clients, @timeout)
        rescue IOError, Errno::EBADF, TypeError
          # Looks like we hit a bad file descriptor or closed connection.
          Moped.logger.debug "replica_set: io error, retrying"
          retry
        end

        return unless readable || errors

        errors.each do |client|
          client.close
          @clients.delete client
        end

        clients, servers = readable.partition { |s| s.class == TCPSocket }

        servers.each do |server|
          Moped.logger.debug "replica_set: accepting new client for #{server.port}"
          @clients << server.accept
        end

        Moped.logger.debug "replica_set: closing dead clients"
        closed, open = clients.partition &:eof?
        closed.each { |client| @clients.delete client }

        if client = open.shift
          Moped.logger.debug "replica_set: finding server for client"
          server = lookup_server(client)

          Moped.logger.debug "replica_set: sending client #{client.inspect} to #{server.port}"
          return server, client
        else
          nil
        end
      end

      def close_clients_for(server)
        Moped.logger.debug "replica_set: closing open clients on #{server.port}"

        @clients.reject! do |client|
          port = client.addr(false)[1]

          if port == server.port
            client.close
            true
          else
            false
          end
        end
      end

      def lookup_server(client)
        port = client.addr(false)[1]

        @servers.find do |server|
          server.to_io && server.to_io.addr[1] == port
        end
      end

    end

  end
end
module Support

  # Module for recording operations.
  #
  #   Support::Stats.install!
  #
  #   stats = Support::Stats.collect do
  #     session.with(safe: true)[:users].insert({})
  #   end
  #
  #   ops = stats["127.0.0.1:27017"]
  #   ops.size # => 2
  #   ops[0].class # => Moped::Protocol::Insert
  #   ops[1].class # => Moped::Protocol::Command
  #
  module Stats
    extend self

    def record(node, operations)
      key = if node.primary?
        :primary
      elsif node.secondary?
        :secondary
      else
        :other
      end

      @stats[key].concat(operations) if @stats
    end

    def collect
      @stats = Hash.new { |hash, key| hash[key] = [] }
      yield
      @stats
    ensure
      @stats = nil
    end

    def install!
      Moped::Node.class_eval <<-EOS
        alias _logging logging

        def logging(operations, &block)
          Support::Stats.record(self, operations)
          _logging(operations, &block)
        end
      EOS
    end

  end
end