From 7d5ab886ed18b397f95235d08924f8b549cd7419 Mon Sep 17 00:00:00 2001
From: Tyler Brock <tyler.brock@gmail.com>
Date: Sun, 29 Jan 2012 19:05:17 -0500
Subject: [PATCH] minor: fix for replica set tests -- each TestCase class now
 gets a new replica set

---
 test/replica_sets/basic_test.rb               | 47 ++++++------
 test/replica_sets/complex_connect_test.rb     | 17 +++--
 test/replica_sets/connect_test.rb             | 46 ++++++------
 test/replica_sets/count_test.rb               | 10 +--
 test/replica_sets/insert_test.rb              | 12 ++--
 test/replica_sets/pooled_insert_test.rb       | 12 ++--
 test/replica_sets/query_test.rb               |  8 +--
 test/replica_sets/read_preference_test.rb     | 22 +++---
 test/replica_sets/refresh_test.rb             | 72 +++++++++----------
 .../replica_sets/refresh_with_threads_test.rb | 12 ++--
 test/replica_sets/replication_ack_test.rb     |  6 +-
 test/replica_sets/rs_test_helper.rb           | 20 +++---
 12 files changed, 145 insertions(+), 139 deletions(-)

diff --git a/test/replica_sets/basic_test.rb b/test/replica_sets/basic_test.rb
index fc64d70..938a7a4 100644
--- a/test/replica_sets/basic_test.rb
+++ b/test/replica_sets/basic_test.rb
@@ -2,50 +2,53 @@ $:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
 require './test/replica_sets/rs_test_helper'
 
 class BasicTest < Test::Unit::TestCase
-  include ReplicaSetTest
+  
+  def setup
+    ensure_rs
+  end
 
   def teardown
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
     @conn.close if defined?(@conn) && @conn
   end
 
   def test_connect
-    @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[1]], [self.rs.host, self.rs.ports[0]],
-      [self.rs.host, self.rs.ports[2]], :name => self.rs.name)
+    @conn = ReplSetConnection.new([@rs.host, @rs.ports[1]], [@rs.host, @rs.ports[0]],
+      [@rs.host, @rs.ports[2]], :name => @rs.name)
     assert @conn.connected?
 
-    assert_equal self.rs.primary, @conn.primary
-    assert_equal self.rs.secondaries.sort, @conn.secondaries.sort
-    assert_equal self.rs.arbiters.sort, @conn.arbiters.sort
+    assert_equal @rs.primary, @conn.primary
+    assert_equal @rs.secondaries.sort, @conn.secondaries.sort
+    assert_equal @rs.arbiters.sort, @conn.arbiters.sort
 
-    @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[1]], [self.rs.host, self.rs.ports[0]],
-      :name => self.rs.name)
+    @conn = ReplSetConnection.new([@rs.host, @rs.ports[1]], [@rs.host, @rs.ports[0]],
+      :name => @rs.name)
     assert @conn.connected?
   end
 
   def test_cache_original_seed_nodes
-    @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[1]], [self.rs.host, self.rs.ports[0]],
-      [self.rs.host, self.rs.ports[2]], [self.rs.host, 19356], :name => self.rs.name)
+    @conn = ReplSetConnection.new([@rs.host, @rs.ports[1]], [@rs.host, @rs.ports[0]],
+      [@rs.host, @rs.ports[2]], [@rs.host, 19356], :name => @rs.name)
     assert @conn.connected?
-    assert @conn.seeds.include?([self.rs.host, 19356]), "Original seed nodes not cached!"
-    assert_equal [self.rs.host, 19356], @conn.seeds.last, "Original seed nodes not cached!"
+    assert @conn.seeds.include?([@rs.host, 19356]), "Original seed nodes not cached!"
+    assert_equal [@rs.host, 19356], @conn.seeds.last, "Original seed nodes not cached!"
   end
 
   def test_accessors
-    seeds = [[self.rs.host, self.rs.ports[0]], [self.rs.host, self.rs.ports[1]],
-      [self.rs.host, self.rs.ports[2]]]
-    args = seeds << {:name => self.rs.name}
+    seeds = [[@rs.host, @rs.ports[0]], [@rs.host, @rs.ports[1]],
+      [@rs.host, @rs.ports[2]]]
+    args = seeds << {:name => @rs.name}
     @conn = ReplSetConnection.new(*args)
 
-    assert_equal @conn.host, self.rs.primary[0]
-    assert_equal @conn.port, self.rs.primary[1]
+    assert_equal @conn.host, @rs.primary[0]
+    assert_equal @conn.port, @rs.primary[1]
     assert_equal @conn.host, @conn.primary_pool.host
     assert_equal @conn.port, @conn.primary_pool.port
     assert_equal @conn.nodes.sort, @conn.seeds.sort
     assert_equal 2, @conn.secondaries.length
     assert_equal 0, @conn.arbiters.length
     assert_equal 2, @conn.secondary_pools.length
-    assert_equal self.rs.name, @conn.replica_set_name
+    assert_equal @rs.name, @conn.replica_set_name
     assert @conn.secondary_pools.include?(@conn.read_pool)
     assert_equal 5, @conn.tag_map.keys.length
     assert_equal 90, @conn.refresh_interval
@@ -55,9 +58,9 @@ class BasicTest < Test::Unit::TestCase
   context "Socket pools" do
     context "checking out writers" do
       setup do
-        seeds = [[self.rs.host, self.rs.ports[0]], [self.rs.host, self.rs.ports[1]],
-          [self.rs.host, self.rs.ports[2]]]
-        args = seeds << {:name => self.rs.name}
+        seeds = [[@rs.host, @rs.ports[0]], [@rs.host, @rs.ports[1]],
+          [@rs.host, @rs.ports[2]]]
+        args = seeds << {:name => @rs.name}
         @con = ReplSetConnection.new(*args)
         @coll = @con[MONGO_TEST_DB]['test-connection-exceptions']
       end
diff --git a/test/replica_sets/complex_connect_test.rb b/test/replica_sets/complex_connect_test.rb
index 145ceeb..6b8b6f4 100644
--- a/test/replica_sets/complex_connect_test.rb
+++ b/test/replica_sets/complex_connect_test.rb
@@ -3,19 +3,22 @@ require 'logger'
 require './test/replica_sets/rs_test_helper'
 
 class ComplexConnectTest < Test::Unit::TestCase
-  include ReplicaSetTest
+  
+  def setup
+    ensure_rs
+  end
 
   def teardown
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
     @conn.close if defined?(@conn) && @conn
   end
 
   def test_complex_connect
     logger = Logger.new(STDOUT)
-    primary = Connection.new(self.rs.host, self.rs.ports[0])
+    primary = Connection.new(@rs.host, @rs.ports[0])
 
-    @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[2]], [self.rs.host, self.rs.ports[1]],
-      [self.rs.host, self.rs.ports[0]], :logger => logger)
+    @conn = ReplSetConnection.new([@rs.host, @rs.ports[2]], [@rs.host, @rs.ports[1]],
+      [@rs.host, @rs.ports[0]], :logger => logger)
 
     @conn['test']['foo'].insert({:a => 1})
     assert @conn['test']['foo'].find_one
@@ -23,13 +26,13 @@ class ComplexConnectTest < Test::Unit::TestCase
     config = primary['local']['system.replset'].find_one
     config['version'] += 1
     config['members'].delete_if do |member|
-      member['host'].include?(self.rs.ports[2].to_s)
+      member['host'].include?(@rs.ports[2].to_s)
     end
 
     assert_raise ConnectionFailure do
       primary['admin'].command({:replSetReconfig => config})
     end
-    self.rs.ensure_up
+    @rs.ensure_up
     assert_raise ConnectionFailure do
       primary['admin'].command({:replSetStepDown => 1})
     end
diff --git a/test/replica_sets/connect_test.rb b/test/replica_sets/connect_test.rb
index c403811..45c2290 100644
--- a/test/replica_sets/connect_test.rb
+++ b/test/replica_sets/connect_test.rb
@@ -2,68 +2,70 @@ $:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
 require './test/replica_sets/rs_test_helper'
 
 class ConnectTest < Test::Unit::TestCase
-  include ReplicaSetTest
+  def setup
+    ensure_rs
+  end
 
   def teardown
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
     @conn.close if defined?(@conn) && @conn
   end
 
   # TODO: test connect timeout.
 
   def test_connect_with_deprecated_multi
-    @conn = Connection.multi([[self.rs.host, self.rs.ports[0]], [self.rs.host, self.rs.ports[1]]], :name => self.rs.name)
+    @conn = Connection.multi([[@rs.host, @rs.ports[0]], [@rs.host, @rs.ports[1]]], :name => @rs.name)
     assert @conn.is_a?(ReplSetConnection)
     assert @conn.connected?
   end
 
   def test_connect_bad_name
     assert_raise_error(ReplicaSetConnectionError, "-wrong") do
-      @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]], [self.rs.host, self.rs.ports[1]],
-        [self.rs.host, self.rs.ports[2]], :name => self.rs.name + "-wrong")
+      @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]], [@rs.host, @rs.ports[1]],
+        [@rs.host, @rs.ports[2]], :name => @rs.name + "-wrong")
     end
   end
 
   def test_connect_with_primary_node_killed
-    node = self.rs.kill_primary
+    node = @rs.kill_primary
 
     # Becuase we're killing the primary and trying to connect right away,
     # this is going to fail right away.
     assert_raise_error(ConnectionFailure, "Failed to connect to primary node") do
-      @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]], [self.rs.host, self.rs.ports[1]],
-        [self.rs.host, self.rs.ports[2]])
+      @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]], [@rs.host, @rs.ports[1]],
+        [@rs.host, @rs.ports[2]])
     end
 
     # This allows the secondary to come up as a primary
     rescue_connection_failure do
-      @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]], [self.rs.host, self.rs.ports[1]],
-        [self.rs.host, self.rs.ports[2]])
+      @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]], [@rs.host, @rs.ports[1]],
+        [@rs.host, @rs.ports[2]])
     end
   end
 
   def test_connect_with_secondary_node_killed
-    node = self.rs.kill_secondary
+    node = @rs.kill_secondary
 
     rescue_connection_failure do
-      @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]], [self.rs.host, self.rs.ports[1]],
-        [self.rs.host, self.rs.ports[2]])
+      @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]], [@rs.host, @rs.ports[1]],
+        [@rs.host, @rs.ports[2]])
     end
     assert @conn.connected?
   end
 
   def test_connect_with_third_node_killed
-    self.rs.kill(self.rs.get_node_from_port(self.rs.ports[2]))
+    @rs.kill(@rs.get_node_from_port(@rs.ports[2]))
 
     rescue_connection_failure do
-      @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]], [self.rs.host, self.rs.ports[1]],
-        [self.rs.host, self.rs.ports[2]])
+      @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]], [@rs.host, @rs.ports[1]],
+        [@rs.host, @rs.ports[2]])
     end
     assert @conn.connected?
   end
 
   def test_connect_with_primary_stepped_down
-    @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]], [self.rs.host, self.rs.ports[1]],
-      [self.rs.host, self.rs.ports[2]])
+    @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]], [@rs.host, @rs.ports[1]],
+      [@rs.host, @rs.ports[2]])
     @conn[MONGO_TEST_DB]['bar'].save({:a => 1}, {:safe => {:w => 3}})
     assert @conn[MONGO_TEST_DB]['bar'].find_one
 
@@ -83,8 +85,8 @@ class ConnectTest < Test::Unit::TestCase
   end
 
   def test_save_with_primary_stepped_down
-    @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]], [self.rs.host, self.rs.ports[1]],
-      [self.rs.host, self.rs.ports[2]])
+    @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]], [@rs.host, @rs.ports[1]],
+      [@rs.host, @rs.ports[2]])
 
     primary = Mongo::Connection.new(@conn.primary_pool.host, @conn.primary_pool.port)
 
@@ -102,13 +104,13 @@ class ConnectTest < Test::Unit::TestCase
   end
 
   def test_connect_with_connection_string
-    @conn = Connection.from_uri("mongodb://#{self.rs.host}:#{self.rs.ports[0]},#{self.rs.host}:#{self.rs.ports[1]}?replicaset=#{self.rs.name}")
+    @conn = Connection.from_uri("mongodb://#{@rs.host}:#{@rs.ports[0]},#{@rs.host}:#{@rs.ports[1]}?replicaset=#{@rs.name}")
     assert @conn.is_a?(ReplSetConnection)
     assert @conn.connected?
   end
 
   def test_connect_with_full_connection_string
-    @conn = Connection.from_uri("mongodb://#{self.rs.host}:#{self.rs.ports[0]},#{self.rs.host}:#{self.rs.ports[1]}?replicaset=#{self.rs.name};safe=true;w=2;fsync=true;slaveok=true")
+    @conn = Connection.from_uri("mongodb://#{@rs.host}:#{@rs.ports[0]},#{@rs.host}:#{@rs.ports[1]}?replicaset=#{@rs.name};safe=true;w=2;fsync=true;slaveok=true")
     assert @conn.is_a?(ReplSetConnection)
     assert @conn.connected?
     assert_equal 2, @conn.safe[:w]
diff --git a/test/replica_sets/count_test.rb b/test/replica_sets/count_test.rb
index 963c215..cf810b9 100644
--- a/test/replica_sets/count_test.rb
+++ b/test/replica_sets/count_test.rb
@@ -2,11 +2,11 @@ $:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
 require './test/replica_sets/rs_test_helper'
 
 class ReplicaSetCountTest < Test::Unit::TestCase
-  include ReplicaSetTest
 
   def setup
-    @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]],
-                                  [self.rs.host, self.rs.ports[1]], [self.rs.host, self.rs.ports[2]],
+    ensure_rs
+    @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]],
+                                  [@rs.host, @rs.ports[1]], [@rs.host, @rs.ports[2]],
                                   :read => :secondary)
     assert @conn.primary_pool
     @primary = Connection.new(@conn.primary_pool.host, @conn.primary_pool.port)
@@ -16,7 +16,7 @@ class ReplicaSetCountTest < Test::Unit::TestCase
   end
 
   def teardown
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
     @conn.close if @conn
   end
 
@@ -25,7 +25,7 @@ class ReplicaSetCountTest < Test::Unit::TestCase
     assert_equal 1, @coll.count
 
     # Kill the current master node
-    @node = self.rs.kill_primary
+    @node = @rs.kill_primary
 
     rescue_connection_failure do
       @coll.insert({:a => 30}, :safe => true)
diff --git a/test/replica_sets/insert_test.rb b/test/replica_sets/insert_test.rb
index 9c89e5f..c009449 100644
--- a/test/replica_sets/insert_test.rb
+++ b/test/replica_sets/insert_test.rb
@@ -2,25 +2,25 @@ $:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
 require './test/replica_sets/rs_test_helper'
 
 class ReplicaSetInsertTest < Test::Unit::TestCase
-  include ReplicaSetTest
 
   def setup
-    @conn = ReplSetConnection.new([TEST_HOST, self.rs.ports[0]],
-                                  [TEST_HOST, self.rs.ports[1]], [TEST_HOST, self.rs.ports[2]])
+    ensure_rs
+    @conn = ReplSetConnection.new([TEST_HOST, @rs.ports[0]],
+                                  [TEST_HOST, @rs.ports[1]], [TEST_HOST, @rs.ports[2]])
     @db = @conn.db(MONGO_TEST_DB)
     @db.drop_collection("test-sets")
     @coll = @db.collection("test-sets")
   end
 
   def teardown
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
     @conn.close if @conn
   end
 
   def test_insert
     @coll.save({:a => 20}, :safe => true)
 
-    self.rs.kill_primary
+    @rs.kill_primary
 
     rescue_connection_failure do
       @coll.save({:a => 30}, :safe => true)
@@ -32,7 +32,7 @@ class ReplicaSetInsertTest < Test::Unit::TestCase
     @coll.save({:a => 70}, :safe => true)
 
     # Restart the old master and wait for sync
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
     sleep(1)
     results = []
 
diff --git a/test/replica_sets/pooled_insert_test.rb b/test/replica_sets/pooled_insert_test.rb
index 3886feb..1fbb733 100644
--- a/test/replica_sets/pooled_insert_test.rb
+++ b/test/replica_sets/pooled_insert_test.rb
@@ -4,18 +4,18 @@ require './test/replica_sets/rs_test_helper'
 # NOTE: This test expects a replica set of three nodes to be running
 # on the local host.
 class ReplicaSetPooledInsertTest < Test::Unit::TestCase
-  include ReplicaSetTest
 
   def setup
-    @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]], [self.rs.host, self.rs.ports[1]],
-      [self.rs.host, self.rs.ports[2]], :pool_size => 5, :timeout => 5, :refresh_mode => false)
+    ensure_rs
+    @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]], [@rs.host, @rs.ports[1]],
+      [@rs.host, @rs.ports[2]], :pool_size => 5, :timeout => 5, :refresh_mode => false)
     @db = @conn.db(MONGO_TEST_DB)
     @db.drop_collection("test-sets")
     @coll = @db.collection("test-sets")
   end
 
   def teardown
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
     @conn.close if @conn
   end
 
@@ -23,7 +23,7 @@ class ReplicaSetPooledInsertTest < Test::Unit::TestCase
     expected_results = [-1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
     @coll.save({:a => -1}, :safe => true)
 
-    self.rs.kill_primary
+    @rs.kill_primary
 
     threads = []
     10.times do |i|
@@ -37,7 +37,7 @@ class ReplicaSetPooledInsertTest < Test::Unit::TestCase
     threads.each {|t| t.join}
 
     # Restart the old master and wait for sync
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
     sleep(1)
     results = []
 
diff --git a/test/replica_sets/query_test.rb b/test/replica_sets/query_test.rb
index 878de1e..3e3988b 100644
--- a/test/replica_sets/query_test.rb
+++ b/test/replica_sets/query_test.rb
@@ -2,17 +2,17 @@ $:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
 require './test/replica_sets/rs_test_helper'
 
 class ReplicaSetQueryTest < Test::Unit::TestCase
-  include ReplicaSetTest
 
   def setup
-    @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]])
+    ensure_rs
+    @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]])
     @db = @conn.db(MONGO_TEST_DB)
     @db.drop_collection("test-sets")
     @coll = @db.collection("test-sets")
   end
 
   def teardown
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
     @conn.close if @conn
   end
 
@@ -28,7 +28,7 @@ class ReplicaSetQueryTest < Test::Unit::TestCase
 
     puts "Benchmark before failover: #{benchmark_queries}"
 
-    self.rs.kill_primary
+    @rs.kill_primary
 
     results = []
     rescue_connection_failure do
diff --git a/test/replica_sets/read_preference_test.rb b/test/replica_sets/read_preference_test.rb
index 2502bdc..74fddc7 100644
--- a/test/replica_sets/read_preference_test.rb
+++ b/test/replica_sets/read_preference_test.rb
@@ -3,12 +3,12 @@ require './test/replica_sets/rs_test_helper'
 require 'logger'
 
 class ReadPreferenceTest < Test::Unit::TestCase
-  include ReplicaSetTest
 
   def setup
+    ensure_rs
     log = Logger.new("test.log")
-    @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]],
-                                 [self.rs.host, self.rs.ports[1]],
+    @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]],
+                                 [@rs.host, @rs.ports[1]],
           :read => :secondary, :pool_size => 50,
           :refresh_mode => false, :refresh_interval => 5, :logger => log)
     @db = @conn.db(MONGO_TEST_DB)
@@ -17,7 +17,7 @@ class ReadPreferenceTest < Test::Unit::TestCase
   end
 
   def teardown
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
   end
 
   def test_read_primary
@@ -35,7 +35,7 @@ class ReadPreferenceTest < Test::Unit::TestCase
   end
 
   def test_query_secondaries
-    @secondary = Connection.new(self.rs.host, @conn.read_pool.port, :slave_ok => true)
+    @secondary = Connection.new(@rs.host, @conn.read_pool.port, :slave_ok => true)
     @coll = @db.collection("test-sets", :safe => {:w => 3, :wtimeout => 20000})
     @coll.save({:a => 20})
     @coll.save({:a => 30})
@@ -49,7 +49,7 @@ class ReadPreferenceTest < Test::Unit::TestCase
     assert results.include?(30)
     assert results.include?(40)
 
-    self.rs.kill_primary
+    @rs.kill_primary
 
     results = []
     rescue_connection_failure do
@@ -68,13 +68,13 @@ class ReadPreferenceTest < Test::Unit::TestCase
     assert_equal 2, @coll.find.to_a.length
 
     # Should still be able to read immediately after killing master node
-    self.rs.kill_primary
+    @rs.kill_primary
     assert_equal 2, @coll.find.to_a.length
     rescue_connection_failure do
       puts "@coll.save()"
       @coll.save({:a => 50}, :safe => {:w => 2, :wtimeout => 10000})
     end
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
     @coll.save({:a => 50}, :safe => {:w => 2, :wtimeout => 10000})
     assert_equal 4, @coll.find.to_a.length
   end
@@ -85,8 +85,8 @@ class ReadPreferenceTest < Test::Unit::TestCase
     @coll.save({:a => 30})
     assert_equal 2, @coll.find.to_a.length
 
-    read_node = self.rs.get_node_from_port(@conn.read_pool.port)
-    self.rs.kill(read_node)
+    read_node = @rs.get_node_from_port(@conn.read_pool.port)
+    @rs.kill(read_node)
 
     # Should fail immediately on next read
     old_read_pool_port = @conn.read_pool.port
@@ -141,7 +141,7 @@ class ReadPreferenceTest < Test::Unit::TestCase
   # end
 
   #def teardown
-  #  self.rs.restart_killed_nodes
+  #  @rs.restart_killed_nodes
   #end
 
 end
diff --git a/test/replica_sets/refresh_test.rb b/test/replica_sets/refresh_test.rb
index 0e99be2..23f0292 100644
--- a/test/replica_sets/refresh_test.rb
+++ b/test/replica_sets/refresh_test.rb
@@ -3,14 +3,14 @@ require './test/replica_sets/rs_test_helper'
 require 'benchmark'
 
 class ReplicaSetRefreshTest < Test::Unit::TestCase
-  include ReplicaSetTest
 
   def setup
+    ensure_rs
     @conn = nil
   end
 
   def teardown
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
     @conn.close if @conn
   end
 
@@ -18,16 +18,16 @@ class ReplicaSetRefreshTest < Test::Unit::TestCase
     Benchmark.bm do |x|
       x.report("Connect") do
         10.times do
-          ReplSetConnection.new([self.rs.host, self.rs.ports[0]],
-                                [self.rs.host, self.rs.ports[1]],
-                                [self.rs.host, self.rs.ports[2]],
+          ReplSetConnection.new([@rs.host, @rs.ports[0]],
+                                [@rs.host, @rs.ports[1]],
+                                [@rs.host, @rs.ports[2]],
                                 :refresh_mode => false)
         end
       end
 
-          @con = ReplSetConnection.new([self.rs.host, self.rs.ports[0]],
-                                       [self.rs.host, self.rs.ports[1]],
-                                       [self.rs.host, self.rs.ports[2]],
+          @con = ReplSetConnection.new([@rs.host, @rs.ports[0]],
+                                       [@rs.host, @rs.ports[1]],
+                                       [@rs.host, @rs.ports[2]],
                                        :refresh_mode => false)
 
       x.report("manager") do
@@ -40,12 +40,12 @@ class ReplicaSetRefreshTest < Test::Unit::TestCase
   end
 
   def test_connect_and_manual_refresh_with_secondaries_down
-    self.rs.kill_all_secondaries
+    @rs.kill_all_secondaries
 
     rescue_connection_failure do
-      @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]],
-                                    [self.rs.host, self.rs.ports[1]],
-                                    [self.rs.host, self.rs.ports[2]],
+      @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]],
+                                    [@rs.host, @rs.ports[1]],
+                                    [@rs.host, @rs.ports[2]],
                                     :refresh_mode => false)
     end
 
@@ -59,7 +59,7 @@ class ReplicaSetRefreshTest < Test::Unit::TestCase
     assert @conn.connected?
     assert_equal @conn.read_pool, @conn.primary_pool
 
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
     assert_equal [], @conn.secondaries
     assert @conn.connected?
     assert_equal @conn.read_pool, @conn.primary_pool
@@ -71,12 +71,12 @@ class ReplicaSetRefreshTest < Test::Unit::TestCase
   end
 
   def test_automated_refresh_with_secondaries_down
-    self.rs.kill_all_secondaries
+    @rs.kill_all_secondaries
 
     rescue_connection_failure do
-      @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]],
-                                    [self.rs.host, self.rs.ports[1]],
-                                    [self.rs.host, self.rs.ports[2]],
+      @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]],
+                                    [@rs.host, @rs.ports[1]],
+                                    [@rs.host, @rs.ports[2]],
                                     :refresh_interval => 2,
                                     :refresh_mode => :sync)
     end
@@ -86,7 +86,7 @@ class ReplicaSetRefreshTest < Test::Unit::TestCase
     assert_equal @conn.read_pool, @conn.primary_pool
     old_refresh_version = @conn.refresh_version
 
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
     sleep(4)
     @conn['foo']['bar'].find_one
     @conn['foo']['bar'].insert({:a => 1})
@@ -100,16 +100,16 @@ class ReplicaSetRefreshTest < Test::Unit::TestCase
   end
   
   def test_automated_refresh_when_secondary_goes_down
-    @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]],
-                                  [self.rs.host, self.rs.ports[1]],
-                                  [self.rs.host, self.rs.ports[2]],
+    @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]],
+                                  [@rs.host, @rs.ports[1]],
+                                  [@rs.host, @rs.ports[2]],
                                   :refresh_interval => 2,
                                   :refresh_mode => :sync)
 
     num_secondaries = @conn.secondary_pools.length
     old_refresh_version = @conn.refresh_version
 
-    n = self.rs.kill_secondary
+    n = @rs.kill_secondary
     sleep(4)
     @conn['foo']['bar'].find_one
 
@@ -118,20 +118,20 @@ class ReplicaSetRefreshTest < Test::Unit::TestCase
     assert_equal num_secondaries - 1, @conn.secondaries.length
     assert_equal num_secondaries - 1, @conn.secondary_pools.length
 
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
   end
 
   def test_automated_refresh_with_removed_node
-    @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]],
-                                  [self.rs.host, self.rs.ports[1]],
-                                  [self.rs.host, self.rs.ports[2]],
+    @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]],
+                                  [@rs.host, @rs.ports[1]],
+                                  [@rs.host, @rs.ports[2]],
                                   :refresh_interval => 2,
                                   :refresh_mode => :sync)
 
     num_secondaries = @conn.secondary_pools.length
     old_refresh_version = @conn.refresh_version
 
-    n = self.rs.remove_secondary_node
+    n = @rs.remove_secondary_node
     sleep(4)
     @conn['foo']['bar'].find_one
 
@@ -140,22 +140,22 @@ class ReplicaSetRefreshTest < Test::Unit::TestCase
     assert_equal num_secondaries - 1, @conn.secondaries.length
     assert_equal num_secondaries - 1, @conn.secondary_pools.length
 
-    self.rs.add_node(n)
+    @rs.add_node(n)
   end
 
   def test_adding_and_removing_nodes
-    @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]],
-                                  [self.rs.host, self.rs.ports[1]],
-                                  [self.rs.host, self.rs.ports[2]],
+    @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]],
+                                  [@rs.host, @rs.ports[1]],
+                                  [@rs.host, @rs.ports[2]],
                                   :refresh_interval => 2, :refresh_mode => :sync)
 
-    self.rs.add_node
+    @rs.add_node
     sleep(4)
     @conn['foo']['bar'].find_one
 
-    @conn2 = ReplSetConnection.new([self.rs.host, self.rs.ports[0]],
-                                   [self.rs.host, self.rs.ports[1]],
-                                   [self.rs.host, self.rs.ports[2]],
+    @conn2 = ReplSetConnection.new([@rs.host, @rs.ports[0]],
+                                   [@rs.host, @rs.ports[1]],
+                                   [@rs.host, @rs.ports[2]],
                                    :refresh_interval => 2, :refresh_mode => :sync)
 
     assert @conn2.secondaries.sort == @conn.secondaries.sort,
@@ -165,7 +165,7 @@ class ReplicaSetRefreshTest < Test::Unit::TestCase
 
     config = @conn['admin'].command({:ismaster => 1})
 
-    self.rs.remove_secondary_node
+    @rs.remove_secondary_node
     sleep(4)
     config = @conn['admin'].command({:ismaster => 1})
 
diff --git a/test/replica_sets/refresh_with_threads_test.rb b/test/replica_sets/refresh_with_threads_test.rb
index f86800d..7c106dc 100644
--- a/test/replica_sets/refresh_with_threads_test.rb
+++ b/test/replica_sets/refresh_with_threads_test.rb
@@ -3,21 +3,21 @@ require './test/replica_sets/rs_test_helper'
 require 'benchmark'
 
 class ReplicaSetRefreshWithThreadsTest < Test::Unit::TestCase
-  include ReplicaSetTest
 
   def setup
+    ensure_rs
     @conn = nil
   end
 
   def teardown
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
     @conn.close if @conn
   end
 
   def test_read_write_load_with_added_nodes
-    @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]],
-                                  [self.rs.host, self.rs.ports[1]],
-                                  [self.rs.host, self.rs.ports[2]],
+    @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]],
+                                  [@rs.host, @rs.ports[1]],
+                                  [@rs.host, @rs.ports[2]],
                                   :refresh_interval => 5,
                                   :refresh_mode => :sync,
                                   :read => :secondary)
@@ -44,7 +44,7 @@ class ReplicaSetRefreshWithThreadsTest < Test::Unit::TestCase
       end
     end
 
-    self.rs.add_node
+    @rs.add_node
     threads.each {|t| t.join }
 
     config = @conn['admin'].command({:ismaster => 1})
diff --git a/test/replica_sets/replication_ack_test.rb b/test/replica_sets/replication_ack_test.rb
index ef6cb63..0feca2c 100644
--- a/test/replica_sets/replication_ack_test.rb
+++ b/test/replica_sets/replication_ack_test.rb
@@ -2,10 +2,10 @@ $:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
 require './test/replica_sets/rs_test_helper'
 
 class ReplicaSetAckTest < Test::Unit::TestCase
-  include ReplicaSetTest
 
   def setup
-    @conn = ReplSetConnection.new([self.rs.host, self.rs.ports[0]])
+    ensure_rs
+    @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]])
 
     @slave1 = Connection.new(@conn.secondary_pools[0].host,
       @conn.secondary_pools[0].port, :slave_ok => true)
@@ -18,7 +18,7 @@ class ReplicaSetAckTest < Test::Unit::TestCase
   end
 
   def teardown
-    self.rs.restart_killed_nodes
+    @rs.restart_killed_nodes
     @conn.close if @conn
   end
 
diff --git a/test/replica_sets/rs_test_helper.rb b/test/replica_sets/rs_test_helper.rb
index 7731cc3..ad390af 100644
--- a/test/replica_sets/rs_test_helper.rb
+++ b/test/replica_sets/rs_test_helper.rb
@@ -2,18 +2,16 @@ $:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
 require './test/test_helper'
 require './test/tools/repl_set_manager'
 
-module ReplicaSetTest
-
-  def self.rs
-    unless defined?(@rs)
-      @rs = ReplSetManager.new
-      @rs.start_set
+class Test::Unit::TestCase
+  # Ensure replica set is available as an instance variable and that
+  # a new set is spun up for each TestCase class
+  def ensure_rs
+    unless defined?(@@current_class) and @@current_class == self.class
+      @@current_class = self.class 
+      @@rs = ReplSetManager.new
+      @@rs.start_set
     end
-    @rs
-  end
-
-  def rs
-    ReplicaSetTest.rs
+    @rs = @@rs
   end
 
   # Generic code for rescuing connection failures and retrying operations.