From c36257f493261bb8802a2d066c8a6579e82f057f Mon Sep 17 00:00:00 2001
From: Michael van Rooijen <michael@vanrooijen.io>
Date: Sun, 9 Jun 2024 22:20:06 +0200
Subject: [PATCH] Add skip_scheduled and skip_working to Resque macro

This allows you to skip the scheduled and working queues when calculating the worker count.
---
 lib/hirefire/macro/resque.rb       | 11 ++++++++++-
 test/hirefire/macro/test_resque.rb | 12 ++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/lib/hirefire/macro/resque.rb b/lib/hirefire/macro/resque.rb
index 94371b8..76949d3 100644
--- a/lib/hirefire/macro/resque.rb
+++ b/lib/hirefire/macro/resque.rb
@@ -21,6 +21,8 @@ module Resque
       #
       # @param queues [Array<String, Symbol>] (optional) Names of the queues for size measurement.
       #   If not provided, size is measured across all queues.
+      # @param [Boolean] skip_scheduled (false) If true, scheduled jobs will not be counted.
+      # @param [Boolean] skip_working (false) If true, working jobs will not be counted.
       # @return [Integer] Total job queue size.
       # @example Calculate size across all queues
       #   HireFire::Macro::Resque.job_queue_size
@@ -28,10 +30,17 @@ module Resque
       #   HireFire::Macro::Resque.job_queue_size(:default)
       # @example Calculate size across the "default" and "mailer" queues
       #   HireFire::Macro::Resque.job_queue_size(:default, :mailer)
-      def job_queue_size(*queues)
+      # @example Calculate size excluding scheduled jobs
+      #   HireFire::Macro::Resque.job_queue_size(skip_scheduled: true)
+      # @example Calculate size excluding jobs currently being worked on
+      #   HireFire::Macro::Resque.job_queue_size(skip_working: true)
+      def job_queue_size(*queues, skip_scheduled: false, skip_working: false)
         queues = normalize_queues(queues, allow_empty: true)
 
         SIZE_METHODS.sum do |size_method|
+          next 0 if skip_scheduled && size_method == :scheduled_size
+          next 0 if skip_working && size_method == :working_size
+
           method(size_method).call(queues)
         end
       end
diff --git a/test/hirefire/macro/test_resque.rb b/test/hirefire/macro/test_resque.rb
index 72225b6..5c0c3a3 100644
--- a/test/hirefire/macro/test_resque.rb
+++ b/test/hirefire/macro/test_resque.rb
@@ -37,8 +37,11 @@ def test_job_queue_size_with_jobs
   def test_job_queue_size_with_working
     enqueue_to_working_with_queue :default, BasicJob
     assert_equal 1, HireFire::Macro::Resque.job_queue_size
+    assert_equal 0, HireFire::Macro::Resque.job_queue_size(skip_working: true)
     assert_equal 1, HireFire::Macro::Resque.job_queue_size(:default)
+    assert_equal 0, HireFire::Macro::Resque.job_queue_size(:default, skip_working: true)
     assert_equal 0, HireFire::Macro::Resque.job_queue_size(:mailer)
+    assert_equal 0, HireFire::Macro::Resque.job_queue_size(:mailer, skip_working: true)
   end
 
   def test_job_queue_size_with_scheduled_jobs
@@ -47,19 +50,28 @@ def test_job_queue_size_with_scheduled_jobs
     Resque.enqueue_in_with_queue(:mailer, 300, BasicJob)
 
     assert_equal 0, HireFire::Macro::Resque.job_queue_size # uncached
+    assert_equal 0, HireFire::Macro::Resque.job_queue_size(skip_scheduled: true) # uncached
 
     Timecop.freeze(Time.now + 200) do
       assert_equal 1, HireFire::Macro::Resque.job_queue_size # uncached
+      assert_equal 0, HireFire::Macro::Resque.job_queue_size(skip_scheduled: true) # uncached
       assert_equal 1, HireFire::Macro::Resque.job_queue_size(:default) # uncached
+      assert_equal 0, HireFire::Macro::Resque.job_queue_size(:default, skip_scheduled: true) # uncached
       assert_equal 0, HireFire::Macro::Resque.job_queue_size(:mailer) # cached
+      assert_equal 0, HireFire::Macro::Resque.job_queue_size(:mailer, skip_scheduled: true) # cached
       assert_equal 1, HireFire::Macro::Resque.job_queue_size(:default, :mailer) # cached
+      assert_equal 0, HireFire::Macro::Resque.job_queue_size(:default, :mailer, skip_scheduled: true) # cached
     end
 
     Timecop.freeze(Time.now + 400) do
       assert_equal 2, HireFire::Macro::Resque.job_queue_size(:default) # expired
+      assert_equal 0, HireFire::Macro::Resque.job_queue_size(:default, skip_scheduled: true) # expired
       assert_equal 1, HireFire::Macro::Resque.job_queue_size(:mailer) # cached
+      assert_equal 0, HireFire::Macro::Resque.job_queue_size(:mailer, skip_scheduled: true) # cached
       assert_equal 3, HireFire::Macro::Resque.job_queue_size(:default, :mailer) # cached
+      assert_equal 0, HireFire::Macro::Resque.job_queue_size(:default, :mailer, skip_scheduled: true) # cached
       assert_equal 3, HireFire::Macro::Resque.job_queue_size # cached
+      assert_equal 0, HireFire::Macro::Resque.job_queue_size(skip_scheduled: true) # cached
     end
   end