From 769cf1413915f987ad2f6c644b27be475cabbbea Mon Sep 17 00:00:00 2001
From: Julik Tarkhanov <me@julik.nl>
Date: Sun, 11 Feb 2024 19:16:36 +0100
Subject: [PATCH] `state` needs to be implemented too

---
 lib/pecorino/cached_throttle.rb | 24 ++++++++++++++++++++++--
 test/cached_throttle_test.rb    | 24 ++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/lib/pecorino/cached_throttle.rb b/lib/pecorino/cached_throttle.rb
index ccc80f2..1ac77e6 100644
--- a/lib/pecorino/cached_throttle.rb
+++ b/lib/pecorino/cached_throttle.rb
@@ -48,6 +48,7 @@ def request(n = 1)
   def able_to_accept?(n = 1)
     blocked_state = cached_blocked_state
     return false if blocked_state&.blocked?
+
     @throttle.able_to_accept?(n)
   end
 
@@ -60,13 +61,32 @@ def throttled(&blk)
     yield
   end
 
+  # Returns the key of the throttle
+  #
+  # @see Pecorino::Throttle#key
+  def key
+    @throttle.key
+  end
+
+  # Returns `false` if there is a currently active block for that throttle in the cache. Otherwise forwards to underlying throttle.
+  #
+  # @see Pecorino::Throttle#able_to_accept?
+  def state
+    blocked_state = cached_blocked_state
+    return blocked_state if blocked_state&.blocked?
+
+    @throttle.state.tap do |state|
+      cache_blocked_state(state) if state.blocked?
+    end
+  end
+
   private
 
   def cache_blocked_state(state)
-    @cache_store.write("pcrn-cached-throttle-#{@throttle.key}", state, expires_after: state.blocked_until)
+    @cache_store.write("pecorino-cached-throttle-state-#{@throttle.key}", state, expires_after: state.blocked_until)
   end
 
   def cached_blocked_state
-    @cache_store.read("pcrn-cached-throttle-#{@throttle.key}")
+    @cache_store.read("pecorino-cached-throttle-state-#{@throttle.key}")
   end
 end
diff --git a/test/cached_throttle_test.rb b/test/cached_throttle_test.rb
index 8ef0db4..7c51f2e 100644
--- a/test/cached_throttle_test.rb
+++ b/test/cached_throttle_test.rb
@@ -94,4 +94,28 @@ class << throttle
 
     assert_nil cached_throttle.throttled { 345 }
   end
+
+  test "caches results of state() and correctly returns cached state until the block is lifted" do
+    store = ActiveSupport::Cache::MemoryStore.new
+    throttle = Pecorino::Throttle.new(key: Random.uuid, capacity: 2, over_time: 1.second, block_for: 2.seconds)
+    cached_throttle = Pecorino::CachedThrottle.new(store, throttle)
+
+    cached_throttle.request(2)
+    blocked_state = cached_throttle.request(1)
+    blocked_state_from_cache = cached_throttle.state
+
+    assert_kind_of Pecorino::Throttle::State, blocked_state
+    assert_kind_of Pecorino::Throttle::State, blocked_state_from_cache
+    assert_predicate blocked_state, :blocked?
+    assert_predicate blocked_state_from_cache, :blocked?
+
+    # Delete the method on the actual throttle as it should not be called anymore until the block is lifted
+    class << throttle
+      undef :state
+    end
+
+    state_from_cache = cached_throttle.state
+    assert_kind_of Pecorino::Throttle::State, state_from_cache
+    assert_predicate state_from_cache, :blocked?
+  end
 end