Log object sizes to a per-user-and-category counter (closes #22)

spec/riak_spec.rb

@@ -1,5 +1,12 @@
 require_relative "spec_helper"
 
+def set_usage_size_info(user, category, size)
+  object = info_bucket.get_or_new("usage:size:#{user}:#{category}")
+  object.content_type = "text/plain"
+  object.data = size.to_s
+  object.store
+end
+
 describe "App with Riak backend" do
   include Rack::Test::Methods
   include RemoteStorage::Riak
@@ -23,6 +30,7 @@ describe "App with Riak backend" do
       last_response.body.must_equal "some text data"
     end
 
+    # If this one fails, try restarting Riak
     it "has a Last-Modified header set" do
       last_response.status.must_equal 200
       last_response.headers["Last-Modified"].wont_be_nil
@@ -85,6 +93,7 @@ describe "App with Riak backend" do
     describe "PUT" do
       before do
         header "Authorization", "Bearer 123"
+        set_usage_size_info "jimmy", "documents", "23"
       end
 
       describe "with implicit content type" do
@@ -102,6 +111,10 @@ describe "App with Riak backend" do
           data_bucket.get("jimmy:documents:bar").content_type.must_equal "text/plain; charset=utf-8"
         end
 
+        it "increases the overall category size" do
+          info_bucket.get("usage:size:jimmy:documents").data.must_equal "35"
+        end
+
         it "indexes the data set" do
           indexes = data_bucket.get("jimmy:documents:bar").indexes
           indexes["user_id_bin"].must_be_kind_of Set
@@ -127,6 +140,10 @@ describe "App with Riak backend" do
           data_bucket.get("jimmy:documents:jason").content_type.must_equal "application/json"
         end
 
+        it "increases the overall category size" do
+          info_bucket.get("usage:size:jimmy:documents").data.must_equal "49"
+        end
+
         it "delivers the data correctly" do
           header "Authorization", "Bearer 123"
           get "/jimmy/documents/jason"
@@ -184,6 +201,40 @@ describe "App with Riak backend" do
         end
       end
 
+      describe "with existing content" do
+        before do
+          set_usage_size_info "jimmy", "documents", "10"
+          put "/jimmy/documents/archive/foo", "lorem ipsum"
+          put "/jimmy/documents/archive/foo", "some awesome content"
+        end
+
+        it "saves the value" do
+          last_response.status.must_equal 200
+          data_bucket.get("jimmy:documents/archive:foo").data.must_equal "some awesome content"
+        end
+
+        it "increases the overall category size" do
+          puts info_bucket.keys.inspect
+          info_bucket.get("usage:size:jimmy:documents").data.must_equal "30"
+        end
+      end
+
+      describe "public data" do
+        before do
+          set_usage_size_info "jimmy", "public/documents", "10"
+          put "/jimmy/public/documents/notes/foo", "note to self"
+        end
+
+        it "saves the value" do
+          last_response.status.must_equal 200
+          data_bucket.get("jimmy:public/documents/notes:foo").data.must_equal "note to self"
+        end
+
+        it "increases the overall category size" do
+          info_bucket.get("usage:size:jimmy:public/documents").data.must_equal "22"
+        end
+      end
+
       context "with binary data" do
         context "binary charset in content-type header" do
           before do
@@ -207,6 +258,10 @@ describe "App with Riak backend" do
             last_response.body.must_equal @image
           end
 
+          it "increases the overall category size" do
+            info_bucket.get("usage:size:jimmy:documents").data.must_equal "16067"
+          end
+
           it "indexes the binary set" do
             indexes = binary_bucket.get("jimmy:documents:jaypeg").indexes
             indexes["user_id_bin"].must_be_kind_of Set
@@ -291,28 +346,33 @@ describe "App with Riak backend" do
     describe "DELETE" do
       before do
         header "Authorization", "Bearer 123"
+        set_usage_size_info "jimmy", "documents", "123"
+        delete "/jimmy/documents/foo"
       end
 
       it "removes the key" do
-        delete "/jimmy/documents/foo"
-
         last_response.status.must_equal 204
         lambda {
           data_bucket.get("jimmy:documents:foo")
         }.must_raise Riak::HTTPFailedRequest
       end
 
+      it "decreases the overall category size" do
+        info_bucket.get("usage:size:jimmy:documents").data.must_equal "101"
+      end
+
       context "binary data" do
         before do
           header "Content-Type", "image/jpeg; charset=binary"
           filename = File.join(File.expand_path(File.dirname(__FILE__)), "fixtures", "rockrule.jpeg")
           @image = File.open(filename, "r").read
           put "/jimmy/documents/jaypeg", @image
+          set_usage_size_info "jimmy", "documents", "100000"
+
+          delete "/jimmy/documents/jaypeg"
         end
 
         it "removes the main object" do
-          delete "/jimmy/documents/jaypeg"
-
           last_response.status.must_equal 204
           lambda {
             data_bucket.get("jimmy:documents:jaypeg")
@@ -320,13 +380,15 @@ describe "App with Riak backend" do
         end
 
         it "removes the binary object" do
-          delete "/jimmy/documents/jaypeg"
-
           last_response.status.must_equal 204
           lambda {
             binary_bucket.get("jimmy:documents:jaypeg")
           }.must_raise Riak::HTTPFailedRequest
         end
+
+        it "decreases the overall category size" do
+          info_bucket.get("usage:size:jimmy:documents").data.must_equal "83956"
+        end
       end
     end
   end

spec/spec_helper.rb

@@ -42,8 +42,12 @@ def binary_bucket
   @binary_bucket ||= storage_client.bucket(settings.bucket_config['binaries'])
 end
 
+def info_bucket
+  @info_bucket ||= storage_client.bucket(settings.bucket_config['info'])
+end
+
 def purge_all_buckets
-  [data_bucket, directory_bucket, auth_bucket, binary_bucket].each do |bucket|
+  [data_bucket, directory_bucket, auth_bucket, binary_bucket, info_bucket].each do |bucket|
     bucket.keys.each {|key| bucket.delete key}
   end
 end