From d93275ab6de1ec06dccc6c3b87659e56256ced5b Mon Sep 17 00:00:00 2001
From: Emily Lin <emily.lin@shopify.com>
Date: Wed, 22 Feb 2023 10:56:35 -0500
Subject: [PATCH] Update total_time for failed jobs

This commit ensures that total_time is updated even when a job
has failed. Previously, total_time would not include the run-time
of the iteration that failed.
---
 lib/job-iteration/iteration.rb |  4 ++--
 test/unit/iteration_test.rb    | 24 ++++++++++++++++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/lib/job-iteration/iteration.rb b/lib/job-iteration/iteration.rb
index ca8c9957..559f93a1 100644
--- a/lib/job-iteration/iteration.rb
+++ b/lib/job-iteration/iteration.rb
@@ -197,9 +197,9 @@ def iterate_with_enumerator(enumerator, arguments)
           "times_interrupted=#{times_interrupted} cursor_position=#{cursor_position}",
       ) unless found_record
 
-      adjust_total_time
-
       true
+    ensure
+      adjust_total_time
     end
 
     def record_unit_of_work(&block)
diff --git a/test/unit/iteration_test.rb b/test/unit/iteration_test.rb
index 275f5ea8..28fbf5ce 100644
--- a/test/unit/iteration_test.rb
+++ b/test/unit/iteration_test.rb
@@ -138,6 +138,21 @@ def each_iteration(*)
     end
   end
 
+  class FailingJob < ActiveJob::Base
+    include JobIteration::Iteration
+    include ActiveSupport::Testing::TimeHelpers
+
+    def build_enumerator(cursor:)
+      enumerator_builder.build_times_enumerator(1, cursor: cursor)
+    end
+
+    def each_iteration(*)
+      travel(10.seconds)
+
+      raise StandardError, "failing job"
+    end
+  end
+
   def test_jobs_that_define_build_enumerator_and_each_iteration_will_not_raise
     push(JobWithRightMethods, "walrus" => "best")
     work_one_job
@@ -352,6 +367,15 @@ def test_max_job_runtime_cannot_be_higher_than_parent
     end
   end
 
+  def test_total_time_is_updated_for_failed_jobs
+    freeze_time do
+      job = FailingJob.new
+      assert_raises(StandardError) { job.perform_now }
+
+      assert_equal(10, job.total_time)
+    end
+  end
+
   private
 
   # Allows building job classes that read max_job_runtime during the test,