pecorino 0.7.1 → 0.7.2

data/test/adapters/adapter_test_methods.rb

@@ -0,0 +1,259 @@
+# The module contains the conformance tests for a storage adapter for Pecorino. A well-behaved adapter
+# should pass all of these tests. When creating a new adapter include this module in your test case
+# and overload the `create_adapter` method
+module AdapterTestMethods
+  LEVEL_DELTA = 0.1
+
+  def adapter
+    @adapter ||= create_adapter
+  end
+
+  def create_adapter
+    raise "Adapter test subclass needs to return an adapter implementation from here."
+  end
+
+  def random_key
+    Random.new(Minitest.seed).hex(4)
+  end
+
+  def test_state_returns_zero_for_nonexistent_bucket
+    k = random_key
+    leak_rate = 2
+    capacity = 3
+
+    level, is_full = adapter.state(key: k, capacity: capacity, leak_rate: leak_rate)
+    assert_equal 0, level
+    assert_equal is_full, false
+  end
+
+  def test_bucket_lifecycle_with_unbounded_fillups
+    k = random_key
+    leak_rate = 2
+    capacity = 1
+
+    level, is_full = adapter.add_tokens(key: k, capacity: capacity, leak_rate: leak_rate, n_tokens: 0.3)
+    assert_in_delta level, 0.3, LEVEL_DELTA
+    assert_equal false, is_full
+
+    level, is_full = adapter.add_tokens(key: k, capacity: capacity, leak_rate: leak_rate, n_tokens: 0.35)
+    assert_in_delta level, 0.65, LEVEL_DELTA
+    assert_equal false, is_full
+
+    level, is_full = adapter.add_tokens(key: k, capacity: capacity, leak_rate: leak_rate, n_tokens: 0.7)
+    assert_in_delta level, 1.0, LEVEL_DELTA
+    assert_equal true, is_full
+
+    level, _ = adapter.state(key: k, capacity: capacity, leak_rate: leak_rate)
+    assert_in_delta level, 1.0, LEVEL_DELTA
+
+    sleep(0.25)
+    level, _ = adapter.state(key: k, capacity: capacity, leak_rate: leak_rate)
+    assert_in_delta level, 0.5, LEVEL_DELTA
+
+    sleep(0.25)
+    level, _ = adapter.state(key: k, capacity: capacity, leak_rate: leak_rate)
+    assert_in_delta level, 0.0, LEVEL_DELTA
+
+    sleep(0.25)
+    level, _ = adapter.state(key: k, capacity: capacity, leak_rate: leak_rate)
+    assert_in_delta level, 0.0, LEVEL_DELTA
+  end
+
+  def test_clamps_fillup_with_negative_value
+    k = random_key
+    leak_rate = 1.1
+    capacity = 15
+
+    level, _, _ = adapter.state(key: k, leak_rate: leak_rate, capacity: capacity)
+    assert_in_delta level, 0, 0.0001
+
+    level, _, _ = adapter.add_tokens(key: k, leak_rate: leak_rate, capacity: capacity, n_tokens: -10)
+    assert_in_delta level, 0, 0.1
+  end
+
+  def test_bucket_lifecycle_with_negative_fillups
+    k = random_key
+    leak_rate = 2
+    capacity = 1
+
+    level, is_full = adapter.add_tokens(key: k, capacity: capacity, leak_rate: leak_rate, n_tokens: 1)
+    assert_in_delta level, 1.0, LEVEL_DELTA
+    assert_equal true, is_full
+
+    level, is_full = adapter.add_tokens(key: k, capacity: capacity, leak_rate: leak_rate, n_tokens: -0.35)
+    assert_in_delta level, 0.65, LEVEL_DELTA
+    assert_equal false, is_full
+
+    level, is_full = adapter.add_tokens(key: k, capacity: capacity, leak_rate: leak_rate, n_tokens: -0.4)
+    assert_in_delta level, 0.25, LEVEL_DELTA
+    assert_equal false, is_full
+
+    level, is_full = adapter.add_tokens(key: k, capacity: capacity, leak_rate: leak_rate, n_tokens: -0.4)
+    assert_in_delta level, 0.0, LEVEL_DELTA
+    assert_equal false, is_full
+  end
+
+  def test_bucket_add_tokens_conditionally_accepts_single_fillup_to_capacity
+    k = random_key
+    leak_rate = 2
+    capacity = 1
+
+    level, is_full, did_accept = adapter.add_tokens_conditionally(key: k, capacity: capacity, leak_rate: leak_rate, n_tokens: 1)
+    assert_in_delta level, 1.0, LEVEL_DELTA
+    assert_equal is_full, true
+    assert_equal did_accept, true
+  end
+
+  def test_bucket_add_tokens_conditionally_accepts_multiple_fillups_to_capacity
+    k = random_key
+    leak_rate = 2
+    capacity = 1
+
+    level, _, did_accept = adapter.add_tokens_conditionally(key: k, capacity: capacity, leak_rate: leak_rate, n_tokens: 0.5)
+    assert_in_delta level, 0.5, LEVEL_DELTA
+    assert_equal did_accept, true
+
+    level, _, did_accept = adapter.add_tokens_conditionally(key: k, capacity: capacity, leak_rate: leak_rate, n_tokens: 0.5)
+    assert_in_delta level, 1.0, LEVEL_DELTA
+    assert_equal did_accept, true
+  end
+
+  def test_bucket_lifecycle_rejects_single_fillup_above_capacity
+    k = random_key
+    leak_rate = 2
+    capacity = 1
+
+    level, is_full, did_accept = adapter.add_tokens_conditionally(key: k, capacity: capacity, leak_rate: leak_rate, n_tokens: 1.2)
+    assert_in_delta level, 0.0, LEVEL_DELTA
+    assert_equal is_full, false
+    assert_equal did_accept, false
+  end
+
+  def test_bucket_lifecycle_rejects_conditional_fillup_that_would_overflow
+    k = random_key
+    leak_rate = 2
+    capacity = 1
+
+    3.times do
+      _level, is_full, did_accept = adapter.add_tokens_conditionally(key: k, capacity: capacity, leak_rate: leak_rate, n_tokens: 0.3)
+      assert_equal is_full, false
+      assert_equal did_accept, true
+    end
+
+    level, is_full, did_accept = adapter.add_tokens_conditionally(key: k, capacity: capacity, leak_rate: leak_rate, n_tokens: 0.3)
+    assert_in_delta level, 0.9, LEVEL_DELTA
+    assert_equal is_full, false
+    assert_equal did_accept, false
+  end
+
+  def test_bucket_lifecycle_handles_conditional_fillup_in_steps
+    key = random_key
+    leak_rate = 1.0
+    capacity = 1.0
+
+    counter = 0
+    try_fillup = ->(fillup_by, should_have_reached_level, should_have_accepted) {
+      counter += 1
+      level, _, did_accept = adapter.add_tokens_conditionally(key: key, capacity: capacity, leak_rate: leak_rate, n_tokens: fillup_by)
+      assert_equal did_accept, should_have_accepted, "Update #{counter} did_accept should be #{should_have_accepted}"
+      assert_in_delta should_have_reached_level, level, 0.1
+    }
+
+    try_fillup.call(1.1, 0.0, false) # Oversized fillup must be refused outright
+    try_fillup.call(0.3, 0.3, true)
+    try_fillup.call(0.3, 0.6, true)
+    try_fillup.call(0.3, 0.9, true)
+    try_fillup.call(0.3, 0.9, false) # Would take the bucket to 1.2, so must be rejected
+
+    sleep(0.2) # Leak out 0.2 tokens
+
+    try_fillup.call(0.3, 1.0, true)
+    try_fillup.call(-2, 0.0, true) # A negative fillup is permitted since it will never take the bucket above capacity
+    try_fillup.call(1.0, 1.0, true) # Filling up in one step should be permitted
+  end
+
+  def test_bucket_lifecycle_allows_conditional_fillup_after_leaking_out
+    rng = Random.new(Minitest.seed)
+    12.times do |i|
+      key = rng.hex(4)
+      capacity = 30
+      leak_rate = capacity / 0.5
+
+      _, _, did_accept = adapter.add_tokens_conditionally(key: key, capacity: capacity, leak_rate: leak_rate, n_tokens: 29.6)
+      assert did_accept, "Should have accepted the topup on iteration #{i}"
+
+      _, _, did_accept = adapter.add_tokens_conditionally(key: key, capacity: capacity, leak_rate: leak_rate, n_tokens: 1)
+      refute did_accept, "Should have refused the topup on iteration #{i}"
+
+      sleep 0.6 # Spend enough time to allow the bucket to leak out completely
+      _, _, did_accept = adapter.add_tokens_conditionally(key: key, capacity: capacity, leak_rate: leak_rate, n_tokens: 1)
+      assert did_accept, "Once the bucket has leaked out to 0 the fillup should be accepted again on iteration #{i}"
+    end
+  end
+
+  def test_set_block_sets_a_block
+    key = random_key
+    now = Time.now.utc
+    block_duration_s = 2.2
+
+    assert_nil adapter.blocked_until(key: key)
+
+    set_block_result = adapter.set_block(key: key, block_for: block_duration_s)
+    assert_kind_of Time, set_block_result
+    assert_in_delta now + block_duration_s, set_block_result, 0.1
+
+    blocked_until = adapter.blocked_until(key: key)
+    assert_in_delta blocked_until, set_block_result, 0.1
+  end
+
+  def test_set_block_does_not_set_block_in_the_past
+    key = random_key
+    assert_nil adapter.blocked_until(key: key)
+    assert_raise(ArgumentError) { adapter.set_block(key: key, block_for: -20) }
+    assert_nil adapter.blocked_until(key: key)
+  end
+
+  def test_set_block_does_not_set_block_which_would_expire_immediately
+    key = random_key
+    assert_nil adapter.blocked_until(key: key)
+    assert_raise(ArgumentError) { adapter.set_block(key: key, block_for: 0) }
+    assert_nil adapter.blocked_until(key: key)
+  end
+
+  def test_prune
+    key = random_key
+    capacity = 30
+    leak_rate = capacity / 0.5
+
+    adapter.add_tokens_conditionally(key: key, capacity: capacity, leak_rate: leak_rate, n_tokens: 29.6)
+    adapter.set_block(key: key, block_for: 0.5)
+
+    sleep 2
+
+    # Both the leaky bucket and the block should have expired by now, and `prune` should not raise
+    assert_nothing_raised { adapter.prune }
+    assert_nil adapter.blocked_until(key: key)
+  end
+
+  def test_create_tables
+    raise "This has to either be a no-op in your test (if your adapter doesn't need any tables) or needs to be written"
+  end
+
+  def xtest_should_accept_threadsafe_conditional_fillups
+    k = random_key
+    capacity = 30
+    leak_rate = capacity / 0.5
+
+    threads = 3.times.map do
+      Thread.new do
+        9.times do
+          adapter.add_tokens_conditionally(key: k, leak_rate: leak_rate, capacity: capacity, n_tokens: 1)
+        end
+      end
+    end
+    threads.map(&:join)
+
+    level, _ = adapter.state(key: k, capacity: capacity, leak_rate: leak_rate)
+    assert_in_delta level, (3 * 9), LEVEL_DELTA
+  end
+end

data/test/adapters/memory_adapter_test.rb

@@ -0,0 +1,12 @@
+require_relative "../test_helper"
+require_relative "adapter_test_methods"
+
+class MemoryAdapterTest < ActiveSupport::TestCase
+  include AdapterTestMethods
+
+  def create_adapter
+    Pecorino::Adapters::MemoryAdapter.new
+  end
+
+  undef :test_create_tables
+end

data/test/adapters/postgres_adapter_test.rb

@@ -0,0 +1,69 @@
+require_relative "../test_helper"
+require_relative "adapter_test_methods"
+
+class PostgresAdapterTest < ActiveSupport::TestCase
+  include AdapterTestMethods
+
+  setup { create_postgres_database_if_none }
+  teardown { truncate_test_tables }
+
+  def self.establish_connection(**options)
+    ActiveRecord::Base.establish_connection(
+      adapter: "postgresql",
+      connect_timeout: 2,
+      **options
+    )
+  end
+
+  def create_adapter
+    Pecorino::Adapters::PostgresAdapter.new(ActiveRecord::Base)
+  end
+
+  SEED_DB_NAME = -> { "pecorino_tests_%s" % Random.new(Minitest.seed).hex(4) }
+
+  def create_postgres_database_if_none
+    self.class.establish_connection(encoding: "unicode", database: SEED_DB_NAME.call)
+    ActiveRecord::Base.connection.execute("SELECT 1 FROM pecorino_leaky_buckets")
+  rescue ActiveRecord::NoDatabaseError, ActiveRecord::ConnectionNotEstablished
+    create_postgres_database
+    retry
+  rescue ActiveRecord::StatementInvalid
+    retained_adapter = adapter # the schema define block is run via instance_exec so it does not retain scope
+    ActiveRecord::Schema.define(version: 1) do |via_definer|
+      retained_adapter.create_tables(via_definer)
+    end
+    retry
+  end
+
+  def create_postgres_database
+    ActiveRecord::Migration.verbose = false
+    self.class.establish_connection(database: "postgres")
+    ActiveRecord::Base.connection.create_database(SEED_DB_NAME.call, charset: :unicode)
+    ActiveRecord::Base.connection.close
+    self.class.establish_connection(encoding: "unicode", database: SEED_DB_NAME.call)
+  end
+
+  def truncate_test_tables
+    ActiveRecord::Base.connection.execute("TRUNCATE TABLE pecorino_leaky_buckets")
+    ActiveRecord::Base.connection.execute("TRUNCATE TABLE pecorino_blocks")
+  end
+
+  def test_create_tables
+    ActiveRecord::Base.transaction do
+      ActiveRecord::Base.connection.execute("DROP TABLE pecorino_leaky_buckets")
+      ActiveRecord::Base.connection.execute("DROP TABLE pecorino_blocks")
+      # The adapter has to be in a variable as the schema definition is scoped to the migrator, not self
+      retained_adapter = create_adapter # the schema define block is run via instance_exec so it does not retain scope
+      ActiveRecord::Schema.define(version: 1) do |via_definer|
+        retained_adapter.create_tables(via_definer)
+      end
+    end
+    assert true
+  end
+
+  Minitest.after_run do
+    ActiveRecord::Base.connection.close
+    establish_connection(database: "postgres")
+    ActiveRecord::Base.connection.drop_database(SEED_DB_NAME.call)
+  end
+end

data/test/adapters/redis_adapter_test.rb

@@ -0,0 +1,27 @@
+require_relative "../test_helper"
+require_relative "adapter_test_methods"
+
+class RedisAdapterTest < ActiveSupport::TestCase
+  include AdapterTestMethods
+
+  teardown { delete_created_keys }
+
+  def create_adapter
+    Pecorino::Adapters::RedisAdapter.new(new_redis, key_prefix: key_prefix)
+  end
+
+  def key_prefix
+    "pecorino-test" + Random.new(Minitest.seed).bytes(4)
+  end
+
+  def delete_created_keys
+    r = new_redis
+    r.del(r.keys(key_prefix + "*"))
+  end
+
+  def new_redis
+    Redis.new
+  end
+
+  undef :test_create_tables
+end

data/test/adapters/sqlite_adapter_test.rb

@@ -0,0 +1,46 @@
+require_relative "../test_helper"
+require_relative "adapter_test_methods"
+
+class SqliteAdapterTest < ActiveSupport::TestCase
+  include AdapterTestMethods
+
+  setup { create_sqlite_db }
+  teardown { drop_sqlite_db }
+
+  def create_sqlite_db
+    ActiveRecord::Migration.verbose = false
+    ActiveRecord::Base.establish_connection(adapter: "sqlite3", database: db_filename)
+
+    ActiveRecord::Schema.define(version: 1) do |via_definer|
+      Pecorino.create_tables(via_definer)
+    end
+  end
+
+  def drop_sqlite_db
+    ActiveRecord::Base.connection.close
+    FileUtils.rm_rf(db_filename)
+    FileUtils.rm_rf(db_filename + "-wal")
+    FileUtils.rm_rf(db_filename + "-shm")
+  end
+
+  def db_filename
+    "pecorino_tests_%s.sqlite3" % Random.new(Minitest.seed).hex(4)
+  end
+
+  def create_adapter
+    Pecorino::Adapters::SqliteAdapter.new(ActiveRecord::Base)
+  end
+
+  def test_create_tables
+    ActiveRecord::Base.transaction do
+      ActiveRecord::Base.connection.execute("DROP TABLE pecorino_leaky_buckets")
+      ActiveRecord::Base.connection.execute("DROP TABLE pecorino_blocks")
+      # The adapter has to be in a variable as the schema definition is scoped to the migrator, not self
+      retained_adapter = create_adapter # the schema define block is run via instance_exec so it does not retain scope
+      ActiveRecord::Schema.define(version: 1) do |via_definer|
+        retained_adapter.create_tables(via_definer)
+      end
+    end
+    assert true
+  end
+end

data/test/block_test.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require "test_helper"
+require "base64"
+
+class BlockTest < ActiveSupport::TestCase
+  test "sets a block" do
+    k = Base64.strict_encode64(Random.bytes(4))
+    assert_nil Pecorino::Block.blocked_until(key: k)
+    assert Pecorino::Block.set!(key: k, block_for: 30.minutes)
+
+    blocked_until = Pecorino::Block.blocked_until(key: k)
+    assert_in_delta Time.now + 30.minutes, blocked_until, 10
+  end
+
+  test "does not return a block which has lapsed" do
+    k = Base64.strict_encode64(Random.bytes(4))
+    assert_nil Pecorino::Block.blocked_until(key: k)
+    Pecorino::Block.set!(key: k, block_for: -30.minutes)
+    blocked_until = Pecorino::Block.blocked_until(key: k)
+    assert_nil blocked_until
+  end
+end

data/test/cached_throttle_test.rb

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+class CachedThrottleTest < ActiveSupport::TestCase
+  def adapter
+    @adapter ||= Pecorino::Adapters::MemoryAdapter.new
+  end
+
+  test "caches results of request! and correctly raises Throttled until the block is lifted" do
+    store = ActiveSupport::Cache::MemoryStore.new
+    throttle = Pecorino::Throttle.new(key: Random.uuid, capacity: 2, over_time: 1, block_for: 10, adapter: adapter)
+    cached_throttle = Pecorino::CachedThrottle.new(store, throttle)
+
+    state1 = cached_throttle.request!
+    state2 = cached_throttle.request!
+    ex = assert_raises(Pecorino::Throttle::Throttled) do
+      cached_throttle.request!
+    end
+
+    assert_kind_of Pecorino::Throttle::State, state1
+    assert_kind_of Pecorino::Throttle::State, state2
+
+    assert_equal throttle, ex.throttle
+
+    # Delete the method on the actual throttle as it should not be called anymore until the block is lifted
+    class << throttle
+      undef :request!
+    end
+    assert_raises(Pecorino::Throttle::Throttled) do
+      cached_throttle.request!
+    end
+  end
+
+  test "caches results of able_to_accept? until the block is lifted" do
+    store = ActiveSupport::Cache::MemoryStore.new
+    throttle = Pecorino::Throttle.new(key: Random.uuid, capacity: 2, over_time: 1, block_for: 10, adapter: adapter)
+    cached_throttle = Pecorino::CachedThrottle.new(store, throttle)
+
+    cached_throttle.request(1)
+    cached_throttle.request(1)
+    cached_throttle.request(1)
+
+    refute cached_throttle.able_to_accept?(1)
+
+    # Delete the method on the actual throttle as it should not be called anymore until the block is lifted
+    class << throttle
+      undef :able_to_accept?
+    end
+
+    refute cached_throttle.able_to_accept?(1)
+  end
+
+  test "caches results of request() 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, block_for: 10, adapter: adapter)
+    cached_throttle = Pecorino::CachedThrottle.new(store, throttle)
+
+    state1 = cached_throttle.request(1)
+    state2 = cached_throttle.request(1)
+    state3 = cached_throttle.request(3)
+
+    assert_kind_of Pecorino::Throttle::State, state1
+    assert_kind_of Pecorino::Throttle::State, state2
+    assert_kind_of Pecorino::Throttle::State, state3
+    assert_predicate state3, :blocked?
+
+    # Delete the method on the actual throttle as it should not be called anymore until the block is lifted
+    class << throttle
+      undef :request
+    end
+    state_from_cache = cached_throttle.request(1)
+    assert_kind_of Pecorino::Throttle::State, state_from_cache
+    assert_predicate state_from_cache, :blocked?
+  end
+
+  test "returns the key of the contained throttle" do
+    store = ActiveSupport::Cache::MemoryStore.new
+    throttle = Pecorino::Throttle.new(key: Random.uuid, capacity: 2, over_time: 1, block_for: 10, adapter: adapter)
+    cached_throttle = Pecorino::CachedThrottle.new(store, throttle)
+    assert_equal cached_throttle.key, throttle.key
+  end
+
+  test "does not run block in throttled() until the block is lifted" do
+    store = ActiveSupport::Cache::MemoryStore.new
+    throttle = Pecorino::Throttle.new(key: Random.uuid, capacity: 2, over_time: 1, block_for: 10, adapter: adapter)
+    cached_throttle = Pecorino::CachedThrottle.new(store, throttle)
+
+    assert_equal 123, cached_throttle.throttled { 123 }
+    assert_equal 234, cached_throttle.throttled { 234 }
+    assert_nil cached_throttle.throttled { 345 }
+
+    # Delete the method on the actual throttle as it should not be called anymore until the block is lifted
+    class << throttle
+      undef :throttled
+    end
+
+    assert_nil cached_throttle.throttled { 345 }
+  end
+end