Add support for pre_count option on Sequel + MySQL

lib/database_cleaner/sequel/truncation.rb

@@ -19,20 +19,30 @@ module DatabaseCleaner
             db.run "TRUNCATE TABLE #{all_tables};"
           end
         else
-          # Truncate each table normally
-          each_table do |db, table|
-            db[table].truncate
+          tables = tables_to_truncate(db)
+
+          if pre_count?
+            # Count rows before truncating
+            pre_count_truncate_tables(db, tables)
+          else
+            # Truncate each table normally
+            truncate_tables(db, tables)
           end
         end
       end
 
-      def each_table
-        tables_to_truncate(db).each do |table|
-          yield db, table.to_sym
-        end
+      private
+
+      def pre_count_truncate_tables(db, tables)
+        tables = tables.reject { |t| db[table.to_sym].count == 0 }
+        truncate_tables(db, tables)
       end
 
-      private
+      def truncate_tables(db, tables)
+        tables.each do |table|
+          db[table.to_sym].truncate
+        end
+      end
 
       def tables_to_truncate(db)
         (@only || db.tables.map(&:to_s)) - @tables_to_exclude
@@ -43,6 +53,10 @@ module DatabaseCleaner
         %w[schema_info schema_migrations]
       end
 
+      def pre_count?
+        @pre_count == true
+      end
+
     end
   end
 end

spec/database_cleaner/sequel/truncation_spec.rb

@@ -130,6 +130,44 @@ module DatabaseCleaner
 
         it_behaves_like 'a Sequel truncation strategy'
         it_behaves_like 'a truncation strategy that resets autoincrement keys by default'
+
+        describe '#pre_count?' do
+          subject { Truncation.new.tap { |t| t.db = db } }
+
+          its(:pre_count?) { should eq false }
+
+          it 'should return true if @reset_id is set and non false or nil' do
+            subject.instance_variable_set(:"@pre_count", true)
+            subject.send(:pre_count?).should eq true
+          end
+
+          it 'should return false if @reset_id is set to false' do
+            subject.instance_variable_set(:"@pre_count", false)
+            subject.send(:pre_count?).should eq false
+          end
+        end
+
+        describe "relying on #pre_count_truncate_tables if asked to" do
+          subject { Truncation.new.tap { |t| t.db = db } }
+
+          it "should rely on #pre_count_truncate_tables if #pre_count? returns true" do
+            subject.instance_variable_set(:"@pre_count", true)
+
+            subject.should_not_receive(:truncate_tables)
+            subject.should_receive(:pre_count_truncate_tables)
+
+            subject.clean
+          end
+
+          it "should not rely on #pre_count_truncate_tables if #pre_count? return false" do
+            subject.instance_variable_set(:"@pre_count", false)
+
+            subject.should_not_receive(:pre_count_truncate_tables)
+            subject.should_receive(:truncate_tables)
+
+            subject.clean
+          end
+        end
       end
     end
     half_supported_configurations.each do |config|