From afb27bc149ee1ed671221e4301757e9a1e0e479e Mon Sep 17 00:00:00 2001
From: Kostas Karachalios <vrinek@me.com>
Date: Fri, 21 Nov 2014 14:42:25 +0100
Subject: [PATCH 1/2] Add support for `pre_count` option on Sequel + MySQL

---
 lib/database_cleaner/sequel/truncation.rb     | 30 +++++++++++----
 .../sequel/truncation_spec.rb                 | 38 +++++++++++++++++++
 2 files changed, 60 insertions(+), 8 deletions(-)

diff --git a/lib/database_cleaner/sequel/truncation.rb b/lib/database_cleaner/sequel/truncation.rb
index b1cfa05..365589f 100755
--- a/lib/database_cleaner/sequel/truncation.rb
+++ b/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
diff --git a/spec/database_cleaner/sequel/truncation_spec.rb b/spec/database_cleaner/sequel/truncation_spec.rb
index 7219776..003201f 100644
--- a/spec/database_cleaner/sequel/truncation_spec.rb
+++ b/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|

From 281ec4f9241ea2e97f05daa0b02fa821e23bc0ac Mon Sep 17 00:00:00 2001
From: Kostas Karachalios <vrinek@me.com>
Date: Mon, 24 Nov 2014 10:26:07 +0100
Subject: [PATCH 2/2] Fix bug in
 DatabaseCleaner::Sequel::Truncation#pre_count_truncate_tables

---
 lib/database_cleaner/sequel/truncation.rb | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/lib/database_cleaner/sequel/truncation.rb b/lib/database_cleaner/sequel/truncation.rb
index 365589f..10f636e 100755
--- a/lib/database_cleaner/sequel/truncation.rb
+++ b/lib/database_cleaner/sequel/truncation.rb
@@ -34,7 +34,7 @@ module DatabaseCleaner
       private
 
       def pre_count_truncate_tables(db, tables)
-        tables = tables.reject { |t| db[table.to_sym].count == 0 }
+        tables = tables.reject { |table| db[table.to_sym].count == 0 }
         truncate_tables(db, tables)
       end
 
@@ -56,9 +56,6 @@ module DatabaseCleaner
       def pre_count?
         @pre_count == true
       end
-
     end
   end
 end
-
-