spec fixes

lib/deimos/active_record_consume/batch_consumption.rb

@@ -159,8 +159,8 @@ def upsert_records(messages)
 
         return if record_list.empty?
 
-        key_col_proc = self.method(:key_columns).to_proc if self.respond_to?(:key_columns)
-        col_proc = self.method(:columns).to_proc if self.respond_to?(:columns)
+        key_col_proc = self.method(:key_columns).to_proc
+        col_proc = self.method(:columns).to_proc
 
         updater = MassUpdater.new(@klass,
                                   key_col_proc: key_col_proc,

lib/deimos/active_record_consume/mass_updater.rb

@@ -7,11 +7,7 @@ class MassUpdater
 
       # @param klass [Class < ActiveRecord::Base]
       def default_keys(klass)
-        klass.connection.
-          indexes(klass.table_name).
-          select(&:unique).
-          map(&:columns).
-          flatten
+        [klass.primary_key]
       end
 
       # @param klass [Class < ActiveRecord::Base]

spec/active_record_batch_consumer_spec.rb

@@ -337,6 +337,9 @@ def deleted_query(records)
         Widget.create!(test_id: 'xxx', some_int: 2, part_one: 'ghi', part_two: 'jkl')
         Widget.create!(test_id: 'yyy', some_int: 7, part_one: 'mno', part_two: 'pqr')
 
+        allow_any_instance_of(MyBatchConsumer).to receive(:key_columns).
+          and_return(%w(part_one part_two))
+
         publish_batch(
           [
             { key: { part_one: 'abc', part_two: 'def' }, # To be created