From ef2694599da4e5e8f33ce3b8d054139fd4b6752b Mon Sep 17 00:00:00 2001
From: Lionel Pereira <lionel.pereira+2021@flipp.com>
Date: Wed, 29 Nov 2023 13:14:18 -0500
Subject: [PATCH] CCOL-2039: Add block to partition! function

---
 .../batch_record_list.rb                      | 24 ++++++++++++-------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/lib/deimos/active_record_consume/batch_record_list.rb b/lib/deimos/active_record_consume/batch_record_list.rb
index f40f0956..c648d51e 100644
--- a/lib/deimos/active_record_consume/batch_record_list.rb
+++ b/lib/deimos/active_record_consume/batch_record_list.rb
@@ -23,19 +23,25 @@ def filter!(method)
         self.batch_records.delete_if { |record| !method.call(record.record) }
       end
 
+      # Filter and return removed invalid batch records by the specified method or block
       # @param method [Proc]
+      # @param block [Proc]
       # @return [Array<BatchRecord>]
-      def partition!(method)
-        valid, invalid = case method.parameters.size
-                         when 2
-          self.batch_records.partition do |record|
-            method.call(record.record, record.associations)
-          end
+      def partition!(method=nil, &block)
+        valid, invalid = if method.nil?
+            self.batch_records.partition(&block)
                          else
-          self.batch_records.partition do |record|
-            method.call(record.record)
+            case method.parameters.size
+            when 2
+              self.batch_records.partition do |record|
+                method.call(record.record, record.associations)
+              end
+            else
+              self.batch_records.partition do |record|
+                method.call(record.record)
+              end
+            end
           end
-                         end
         self.batch_records = valid
         invalid
       end