From 3bf9659d564e193d9b1bcfccf80bec4070e205d6 Mon Sep 17 00:00:00 2001
From: Geremia Taglialatela <tagliala.dev@gmail.com>
Date: Sat, 3 Feb 2024 10:34:51 +0100
Subject: [PATCH] Fix `find_each` and `in_batches` methods

Close #177
---
 .rubocop_todo.yml                              |  2 +-
 lib/chrono_model.rb                            |  2 +-
 lib/chrono_model/patches/batches.rb            | 12 ++++--------
 lib/chrono_model/version.rb                    |  2 +-
 spec/chrono_model/time_machine/batches_spec.rb |  4 ++++
 5 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/.rubocop_todo.yml b/.rubocop_todo.yml
index 315f7c2b..e23a03b2 100644
--- a/.rubocop_todo.yml
+++ b/.rubocop_todo.yml
@@ -1,6 +1,6 @@
 # This configuration was generated by
 # `rubocop --auto-gen-config --no-offense-counts --no-auto-gen-timestamp`
-# using RuboCop version 1.59.0.
+# using RuboCop version 1.60.2.
 # The point is for the user to remove these configuration records
 # one by one as the offenses are removed from the code base.
 # Note that changes in the inspected code, or installation of new
diff --git a/lib/chrono_model.rb b/lib/chrono_model.rb
index a7ce306f..d4a87fd9 100644
--- a/lib/chrono_model.rb
+++ b/lib/chrono_model.rb
@@ -56,7 +56,7 @@ def self.history_models
 
   ActiveRecord::Associations::Preloader::ThroughAssociation.prepend ChronoModel::Patches::Preloader::ThroughAssociation
 
-  ActiveRecord::Batches::BatchEnumerator.prepend ChronoModel::Patches::Batches::BatchEnumerator
+  ActiveRecord::Batches.prepend ChronoModel::Patches::Batches
 end
 
 ActiveSupport.on_load :after_initialize do
diff --git a/lib/chrono_model/patches/batches.rb b/lib/chrono_model/patches/batches.rb
index 5d7e52ca..debacbe5 100644
--- a/lib/chrono_model/patches/batches.rb
+++ b/lib/chrono_model/patches/batches.rb
@@ -3,14 +3,10 @@
 module ChronoModel
   module Patches
     module Batches
-      module BatchEnumerator
-        def each(&block)
-          if @relation.try(:history?)
-            @relation.with_hid_pkey { super }
-          else
-            super
-          end
-        end
+      def in_batches(**)
+        return super unless try(:history?)
+
+        with_hid_pkey { super }
       end
     end
   end
diff --git a/lib/chrono_model/version.rb b/lib/chrono_model/version.rb
index 55ac1f44..33d68ad0 100644
--- a/lib/chrono_model/version.rb
+++ b/lib/chrono_model/version.rb
@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module ChronoModel
-  VERSION = '3.0.0'
+  VERSION = '3.0.1'
 end
diff --git a/spec/chrono_model/time_machine/batches_spec.rb b/spec/chrono_model/time_machine/batches_spec.rb
index b44ceb43..050e2ea3 100644
--- a/spec/chrono_model/time_machine/batches_spec.rb
+++ b/spec/chrono_model/time_machine/batches_spec.rb
@@ -6,6 +6,10 @@
 RSpec.describe ChronoModel::TimeMachine do
   include ChronoTest::TimeMachine::Helpers
 
+  describe '.find_each' do
+    it { expect(Foo.history.find_each(batch_size: 2).count).to eq Foo.history.count }
+  end
+
   describe '.in_batches' do
     let(:foo_in_batches_of_two) do
       [