From c0b4b66cc1a2da566ad0a0dde9e7c4c3622b51c3 Mon Sep 17 00:00:00 2001
From: Rafael Gibim <9031589+Drowze@users.noreply.github.com>
Date: Mon, 4 Sep 2023 17:08:44 -0300
Subject: [PATCH] Add helper to cleanup delayed sidekiq timechunks

Mainly useful to avoid flaky tests when using a real redis database
and not flushing the database in between tests
---
 CHANGELOG.md                                    |  2 ++
 README.md                                       |  6 ++++++
 lib/chewy/strategy/delayed_sidekiq.rb           | 13 +++++++++++++
 lib/chewy/strategy/delayed_sidekiq/scheduler.rb |  3 +++
 spec/chewy/strategy/delayed_sidekiq_spec.rb     | 12 ++++++++++++
 5 files changed, 36 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ee8a9ba74..eb2863c94 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,8 @@
 
 ### New Features
 
+* [#894](https://github.com/toptal/chewy/pull/894): Way of cleaning redis from artifacts left by `delayed_sidekiq` strategy which could potentially cause flaky tests. ([@Drowze](https://github.com/Drowze))
+
 ### Changes
 
 ### Bugs Fixed
diff --git a/README.md b/README.md
index bc08c32a8..dbfe974ce 100644
--- a/README.md
+++ b/README.md
@@ -848,6 +848,12 @@ Explicit call of the reindex using `:delayed_sidekiq` strategy with `:update_fie
 CitiesIndex.import([1, 2, 3], update_fields: [:name], strategy: :delayed_sidekiq)
 ```
 
+While running tests with delayed_sidekiq strategy and Sidekiq is using a real redis instance that is NOT cleaned up in between tests (via e.g. `Sidekiq.redis(&:flushdb)`), you'll want to cleanup some redis keys in between tests to avoid state leaking and flaky tests. Chewy provides a convenience method for that:
+```ruby
+# it might be a good idea to also add to your testing setup, e.g.: a rspec `before` hook
+Chewy::Strategy::DelayedSidekiq.clear_timechunks!
+```
+
 #### `:active_job`
 
 This does the same thing as `:atomic`, but using ActiveJob. This will inherit the ActiveJob configuration settings including the `active_job.queue_adapter` setting for the environment. Patch `Chewy::Strategy::ActiveJob::Worker` for index updates improving.
diff --git a/lib/chewy/strategy/delayed_sidekiq.rb b/lib/chewy/strategy/delayed_sidekiq.rb
index 2c694071c..0bdd25c88 100644
--- a/lib/chewy/strategy/delayed_sidekiq.rb
+++ b/lib/chewy/strategy/delayed_sidekiq.rb
@@ -5,6 +5,19 @@ class Strategy
     class DelayedSidekiq < Sidekiq
       require_relative 'delayed_sidekiq/scheduler'
 
+      # cleanup the redis sets used internally. Useful mainly in tests to avoid
+      # leak and potential flaky tests.
+      def self.clear_timechunks!
+        ::Sidekiq.redis do |redis|
+          timechunk_sets = redis.smembers(Chewy::Strategy::DelayedSidekiq::Scheduler::ALL_SETS_KEY)
+          break if timechunk_sets.empty?
+
+          redis.pipelined do |pipeline|
+            timechunk_sets.each { |set| pipeline.del(set) }
+          end
+        end
+      end
+
       def leave
         @stash.each do |type, ids|
           next if ids.empty?
diff --git a/lib/chewy/strategy/delayed_sidekiq/scheduler.rb b/lib/chewy/strategy/delayed_sidekiq/scheduler.rb
index 9f4bf386e..f1010a3ee 100644
--- a/lib/chewy/strategy/delayed_sidekiq/scheduler.rb
+++ b/lib/chewy/strategy/delayed_sidekiq/scheduler.rb
@@ -18,6 +18,7 @@ class Scheduler
         DEFAULT_MARGIN = 2
         DEFAULT_QUEUE = 'chewy'
         KEY_PREFIX = 'chewy:delayed_sidekiq'
+        ALL_SETS_KEY = "#{KEY_PREFIX}:all_sets".freeze
         FALLBACK_FIELDS = 'all'
         FIELDS_IDS_SEPARATOR = ';'
         IDS_SEPARATOR = ','
@@ -68,8 +69,10 @@ def postpone
           ::Sidekiq.redis do |redis|
             # warning: Redis#sadd will always return an Integer in Redis 5.0.0. Use Redis#sadd? instead
             if redis.respond_to?(:sadd?)
+              redis.sadd?(ALL_SETS_KEY, timechunks_key)
               redis.sadd?(timechunk_key, serialize_data)
             else
+              redis.sadd(ALL_SETS_KEY, timechunks_key)
               redis.sadd(timechunk_key, serialize_data)
             end
 
diff --git a/spec/chewy/strategy/delayed_sidekiq_spec.rb b/spec/chewy/strategy/delayed_sidekiq_spec.rb
index ac9f792c7..c0d43e576 100644
--- a/spec/chewy/strategy/delayed_sidekiq_spec.rb
+++ b/spec/chewy/strategy/delayed_sidekiq_spec.rb
@@ -186,5 +186,17 @@
         end
       end
     end
+
+    describe '#clear_delayed_sidekiq_timechunks test helper' do
+      it 'clears redis from the timechunk sorted sets to avoid leak between tests' do
+        timechunks_set = -> { Sidekiq.redis { |redis| redis.zrange('chewy:delayed_sidekiq:CitiesIndex:timechunks', 0, -1) } }
+
+        expect { CitiesIndex.import!([1], strategy: :delayed_sidekiq) }
+          .to change { timechunks_set.call.size }.by(1)
+
+        expect { Chewy::Strategy::DelayedSidekiq.clear_timechunks! }
+          .to change { timechunks_set.call.size }.to(0)
+      end
+    end
   end
 end