diff --git a/lib/database_cleaner/sequel/truncation.rb b/lib/database_cleaner/sequel/truncation.rb index b1cfa05..10f636e 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 { |table| 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,8 +53,9 @@ 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|