implement token bucket rate limiting

test/functional/posts_controller_test.rb

@@ -19,16 +19,18 @@ class PostsControllerTest < ActionController::TestCase
     context "for api calls" do
       context "passing the api limit" do
         setup do
-          User.any_instance.stubs(:api_hourly_limit).returns(5)
+          @post = FactoryGirl.create(:post)
+          @bucket = TokenBucket.create(user_id: @user.id, token_count: 5, last_touched_at: Time.now)
+          User.any_instance.stubs(:api_burst_limit).returns(5)
         end
         
         should "work" do
-          CurrentUser.user.api_hourly_limit.times do
-            get :index, {:format => "json", :login => @user.name, :api_key => @user.api_key.key}
+          @user.api_burst_limit.times do
+            post :update, {:format => "json", :id => @post.id, :post => {:rating => "q"}, :login => @user.name, :api_key => @user.api_key.key}
             assert_response :success
           end
 
-          get :index, {:format => "json", :login => @user.name, :api_key => @user.api_key.key}
+            post :update, {:format => "json", :id => @post.id, :post => {:rating => "q"}, :login => @user.name, :api_key => @user.api_key.key}
           assert_response 429
         end
       end

test/unit/api_limiter_test.rb

@@ -1,43 +0,0 @@
-require 'test_helper'
-
-class ApiLimiterTest < ActiveSupport::TestCase
-  context "for reads" do
-    context "for an anonymous user" do
-      setup do
-        @count = 5
-        @user = AnonymousUser.new
-        CurrentUser.user = @user
-      end
-
-      should "respect api limits" do
-        @user.expects(:api_hourly_limit).with(true).times(@count + 1).returns(@count)
-
-        @count.times do
-          assert_equal(false, ApiLimiter.throttled?(CurrentUser.id || "127.0.0.1", "GET"))
-        end
-
-        assert_equal(true, ApiLimiter.throttled?(CurrentUser.id || "127.0.0.1", "GET"))
-      end
-    end
-  end
-
-  context "for writes" do
-    context "for an anonymous user" do
-      setup do
-        @count = 5
-        @user = AnonymousUser.new
-        CurrentUser.user = @user
-      end
-
-      should "respect api limits" do
-        @user.expects(:api_hourly_limit).with(false).times(@count + 1).returns(@count)
-
-        @count.times do
-          assert_equal(false, ApiLimiter.throttled?(CurrentUser.id || "127.0.0.1", "POST"))
-        end
-
-        assert_equal(true, ApiLimiter.throttled?(CurrentUser.id || "127.0.0.1", "POST"))
-      end
-    end
-  end
-end

test/unit/token_bucket_test.rb

@@ -0,0 +1,45 @@
+require 'test_helper'
+
+class TokenBucketTest < ActiveSupport::TestCase
+  context "#add!" do
+    setup do
+      @user = FactoryGirl.create(:user)
+      TokenBucket.create(user_id: @user.id, last_touched_at: 1.minute.ago, token_count: 0)
+    end
+
+    should "work" do
+      @user.token_bucket.add!
+      assert_operator(@user.token_bucket.token_count, :>, 0)
+      @user.reload
+      assert_operator(@user.token_bucket.token_count, :>, 0)
+    end
+  end
+
+  context "#consume!" do
+    setup do
+      @user = FactoryGirl.create(:user)
+      TokenBucket.create(user_id: @user.id, last_touched_at: 1.minute.ago, token_count: 1)
+    end
+
+    should "work" do
+      @user.token_bucket.consume!
+      assert_operator(@user.token_bucket.token_count, :<, 1)
+      @user.reload     
+      assert_operator(@user.token_bucket.token_count, :<, 1)
+    end
+  end
+
+  context "#throttled?" do
+    setup do
+      @user = FactoryGirl.create(:user)
+      TokenBucket.create(user_id: @user.id, last_touched_at: 1.minute.ago, token_count: 0)
+    end
+
+    should "work" do
+      assert(!@user.token_bucket.throttled?)
+      assert_operator(@user.token_bucket.token_count, :<, 60)
+      @user.reload
+      assert_operator(@user.token_bucket.token_count, :<, 60)
+    end
+  end
+end