modified truncation to support multiple connections

lib/database_cleaner/active_record/truncation.rb

@@ -72,21 +72,31 @@ module DatabaseCleaner::ActiveRecord
   class Truncation < ::DatabaseCleaner::TruncationBase
 
     def clean
-      connection.disable_referential_integrity do
-        tables_to_truncate.each do |table_name|
-          connection.truncate_table table_name
+      connections.each do |connection|
+        connection.disable_referential_integrity do
+          tables_to_truncate_for_connection(connection).each do |table_name|
+            connection.truncate_table table_name
+          end
         end
       end
     end
 
     private
 
-    def tables_to_truncate
+    def tables_to_truncate_for_connection(connection)
       (@only || connection.tables) - @tables_to_exclude
     end
+    
+    # def tables_to_truncate
+    #   (@only || connection.tables) - @tables_to_exclude
+    # end
 
-    def connection
-      ::ActiveRecord::Base.connection
+    # def connection
+    #   ::ActiveRecord::Base.connection
+    # end
+    
+    def connections
+      DatabaseCleaner::ActiveRecord.connection_klasses.map{ |klass| klass.connection }
     end
 
     # overwritten

spec/database_cleaner/active_record/truncation_spec.rb

@@ -15,51 +15,121 @@ module DatabaseCleaner
   module ActiveRecord
 
     describe Truncation do
+      let (:connection) { mock('connection') }
+      
       before(:each) do
-        @connection = mock('connection')
-        @connection.stub!(:disable_referential_integrity).and_yield
-        ::ActiveRecord::Base.stub!(:connection).and_return(@connection)
+        connection.stub!(:disable_referential_integrity).and_yield
+        ::ActiveRecord::Base.stub!(:connection).and_return(connection)
       end
 
-      it "should truncate all tables except for schema_migrations" do
-        @connection.stub!(:tables).and_return(%w[schema_migrations widgets dogs])
+      context "single connection" do
+        it "should truncate all tables except for schema_migrations" do
+          connection.stub!(:tables).and_return(%w[schema_migrations widgets dogs])
 
-        @connection.should_receive(:truncate_table).with('widgets')
-        @connection.should_receive(:truncate_table).with('dogs')
-        @connection.should_not_receive(:truncate_table).with('schema_migrations')
+          connection.should_receive(:truncate_table).with('widgets')
+          connection.should_receive(:truncate_table).with('dogs')
+          connection.should_not_receive(:truncate_table).with('schema_migrations')
 
-        Truncation.new.clean
+          Truncation.new.clean
+        end
+
+        it "should only truncate the tables specified in the :only option when provided" do
+          connection.stub!(:tables).and_return(%w[schema_migrations widgets dogs])
+
+          connection.should_receive(:truncate_table).with('widgets')
+          connection.should_not_receive(:truncate_table).with('dogs')
+
+          Truncation.new(:only => ['widgets']).clean
+        end
+
+        it "should not truncate the tables specified in the :except option" do
+          connection.stub!(:tables).and_return(%w[schema_migrations widgets dogs])
+
+          connection.should_receive(:truncate_table).with('dogs')
+          connection.should_not_receive(:truncate_table).with('widgets')
+
+          Truncation.new(:except => ['widgets']).clean
+        end
+
+        it "should raise an error when :only and :except options are used" do
+          running {
+            Truncation.new(:except => ['widgets'], :only => ['widgets'])
+          }.should raise_error(ArgumentError)
+        end
+
+        it "should raise an error when invalid options are provided" do
+          running { Truncation.new(:foo => 'bar') }.should raise_error(ArgumentError)
+        end
       end
+      
+      context "multiple connection" do
+        
+        let (:another_connection) { mock('another connection') }
+        let (:model_klass) { mock('klass') }
+        
+        before(:each) do
+          another_connection.stub!(:disable_referential_integrity).and_yield
+          model_klass.stub!(:connection).and_return(another_connection)
+          
+          DatabaseCleaner::ActiveRecord.connection_klasses = [model_klass]
+        end
 
-      it "should only truncate the tables specified in the :only option when provided" do
-        @connection.stub!(:tables).and_return(%w[schema_migrations widgets dogs])
+        after(:each) do
+          DatabaseCleaner::ActiveRecord.connection_klasses = []
+        end
+        
+        it "should truncate all tables except for schema_migrations on both connections" do
+          connection.stub!(:tables).and_return(%w[schema_migrations widgets dogs])
+          another_connection.stub!(:tables).and_return(%w[schema_migrations widgets dogs])
 
-        @connection.should_receive(:truncate_table).with('widgets')
-        @connection.should_not_receive(:truncate_table).with('dogs')
+          connection.should_receive(:truncate_table).with('widgets')
+          connection.should_receive(:truncate_table).with('dogs')
+          connection.should_not_receive(:truncate_table).with('schema_migrations')
+          
+          another_connection.should_receive(:truncate_table).with('widgets')
+          another_connection.should_receive(:truncate_table).with('dogs')
+          another_connection.should_not_receive(:truncate_table).with('schema_migrations')
+          
+          Truncation.new.clean
+        end
 
-        Truncation.new(:only => ['widgets']).clean
+        it "should only truncate the tables specified in the :only option when provided" do
+          connection.stub!(:tables).and_return(%w[schema_migrations widgets dogs])
+          another_connection.stub!(:tables).and_return(%w[schema_migrations widgets dogs])
+
+          connection.should_receive(:truncate_table).with('widgets')
+          connection.should_not_receive(:truncate_table).with('dogs')
+          
+          another_connection.should_receive(:truncate_table).with('widgets')
+          another_connection.should_not_receive(:truncate_table).with('dogs')
+
+          Truncation.new(:only => ['widgets']).clean
+        end
+
+        it "should not truncate the tables specified in the :except option" do
+          connection.stub!(:tables).and_return(%w[schema_migrations widgets dogs])
+          another_connection.stub!(:tables).and_return(%w[schema_migrations widgets dogs])
+
+          connection.should_receive(:truncate_table).with('dogs')
+          connection.should_not_receive(:truncate_table).with('widgets')
+          
+          another_connection.should_receive(:truncate_table).with('dogs')
+          another_connection.should_not_receive(:truncate_table).with('widgets')
+
+          Truncation.new(:except => ['widgets']).clean
+        end
+
+        it "should raise an error when :only and :except options are used" do
+          running {
+            Truncation.new(:except => ['widgets'], :only => ['widgets'])
+          }.should raise_error(ArgumentError)
+        end
+
+        it "should raise an error when invalid options are provided" do
+          running { Truncation.new(:foo => 'bar') }.should raise_error(ArgumentError)
+        end
       end
-
-      it "should not truncate the tables specified in the :except option" do
-        @connection.stub!(:tables).and_return(%w[schema_migrations widgets dogs])
-
-        @connection.should_receive(:truncate_table).with('dogs')
-        @connection.should_not_receive(:truncate_table).with('widgets')
-
-        Truncation.new(:except => ['widgets']).clean
-      end
-
-      it "should raise an error when :only and :except options are used" do
-        running {
-          Truncation.new(:except => ['widgets'], :only => ['widgets'])
-        }.should raise_error(ArgumentError)
-      end
-
-      it "should raise an error when invalid options are provided" do
-        running { Truncation.new(:foo => 'bar') }.should raise_error(ArgumentError)
-      end
-
-
+      
     end
 
   end