Make resettable riak a first class data store.

lib/curator.rb

@@ -5,7 +5,10 @@
 require 'curator/model'
 require 'curator/repository'
 require 'curator/configuration'
+require 'curator/riak/configuration'
 require 'curator/riak/data_store'
+require 'curator/resettable_riak/configuration'
+require 'curator/resettable_riak/data_store'
 require 'curator/railtie' if defined?(Rails)
 
 module Curator
@@ -15,11 +18,14 @@ class << self
 
   def self.configure(data_store, &block)
     path = "curator/#{data_store.to_s}/configuration"
-    require path
-    @config ||= path.camelize.constantize.new
+    @config = path.camelize.constantize.new
     yield(@config) if block_given?
   end
 
+  def self.data_store
+    config.data_store
+  end
+
   self.configure(:riak) do |config|
     config.environment = 'development'
     config.migrations_path = File.expand_path(File.dirname(__FILE__) + "/../db/migrate")

lib/curator/repository.rb

@@ -13,7 +13,7 @@ def collection_name
       end
 
       def data_store
-        @data_store ||= Riak::DataStore
+        @data_store ||= Curator.data_store
       end
 
       def data_store=(store)

lib/curator/resettable_riak/configuration.rb

@@ -0,0 +1,7 @@
+module Curator::ResettableRiak
+  class Configuration < Curator::Riak::Configuration
+    def data_store
+      Curator::ResettableRiak::DataStore
+    end
+  end
+end

lib/curator/resettable_riak/data_store.rb

@@ -0,0 +1,50 @@
+require 'active_support/core_ext/kernel/reporting'
+
+module Curator
+  module ResettableRiak
+    class DataStore < Riak::DataStore
+      def self.bucket_prefix
+        job = "#{ENV['JOB_NAME'].gsub(/[^[:alnum:]]/, '_')}" if ENV['JOB_NAME'].present?
+        [Curator.config.bucket_prefix, job, Curator.config.environment].compact.join(':')
+      end
+
+      def self.exclude_from_reset
+        @exclude_from_reset = true
+        yield
+        @exclude_from_reset = false
+      end
+
+      def self.remove_all_keys
+        silence_warnings do
+          buckets = client.buckets.select { |bucket| bucket.name.start_with?(DataStore.bucket_prefix) }
+          buckets.each do |bucket|
+            bucket.keys do |keys|
+              keys.each { |key| bucket.delete(key) }
+            end
+          end
+        end
+      end
+
+      def self.reset!
+        @bucket_names ||= {}
+        deletable_buckets = @bucket_names.each do |bucket_name, keys|
+          bucket = _bucket(bucket_name)
+          keys.each {|key| bucket.delete(key)}
+        end
+        @bucket_names = {}
+      end
+
+      def self.save(options)
+        result = super
+
+        unless @exclude_from_reset
+          @bucket_names ||= {}
+          @bucket_names[options[:collection_name]] ||= []
+          @bucket_names[options[:collection_name]] << result.key
+        end
+
+        result
+      end
+    end
+  end
+end

lib/curator/riak/configuration.rb

@@ -3,5 +3,9 @@ class Configuration
     include Curator::Configuration
 
     attr_accessor :bucket_prefix, :riak_config_file
+
+    def data_store
+      Curator::Riak::DataStore
+    end
   end
 end

spec/curator/repository_spec.rb

@@ -83,7 +83,7 @@
 
       repository.save(model)
 
-      riak_data = Curator::Riak::Test::ResettableDataStore.find_by_key("test_models", model.id)[:data]
+      riak_data = Curator.data_store.find_by_key("test_models", model.id)[:data]
       riak_data.has_key?("some_field").should be_false
     end
 

spec/curator_spec.rb

@@ -2,24 +2,28 @@
 
 describe Curator do
   describe 'self.configure' do
+    around(:each) do |example|
+      old_config = Curator.config
+      example.run
+      Curator.instance_variable_set('@config', old_config)
+    end
+
     it 'creates a configuration if one does not exist' do
       old_config = Curator.config
       Curator.instance_variable_set(:@config, nil)
-      Curator.configure(:riak)
+      Curator.configure(:resettable_riak)
       Curator.config.should be_kind_of(Curator::Configuration)
-      Curator.instance_variable_set(:@config, old_config)
     end
 
-    it 'leaves existing configuration in place' do
-      Curator.configure(:riak)
-      old_config = Curator.config
-      Curator.configure(:riak)
-      Curator.config.should equal(old_config)
+    it 'overwrites existing configuration in place' do
+      Curator.configure(:resettable_riak) { |config| config.environment = "orig" }
+      Curator.configure(:resettable_riak) { |config| config.environment = "new" }
+      Curator.config.environment.should == "new"
     end
 
     it 'takes a block and passes the configuration' do
       block_config = nil
-      Curator.configure(:riak) do |config|
+      Curator.configure(:resettable_riak) do |config|
         block_config = config
       end
 

spec/spec_helper.rb

@@ -1,21 +1,21 @@
 require 'i18n'
 require 'curator'
 require 'timecop'
-require 'curator/riak/test/cleaner'
-require 'curator/riak/test/resettable_data_store'
 
-Curator.configure(:riak) do |config|
+Curator.configure(:resettable_riak) do |config|
   config.environment = "test"
   config.migrations_path = "/tmp/curator_migrations"
+  config.bucket_prefix = 'curator'
+  config.riak_config_file = File.expand_path(File.dirname(__FILE__) + "/../config/riak.yml")
 end
 
 RSpec.configure do |config|
   config.before(:suite) do
-    Curator::Riak::Test::Cleaner.remove_all_keys
+    Curator.data_store.remove_all_keys
   end
 
   config.after(:each) do
-    Curator::Riak::Test::ResettableDataStore.reset!
+    Curator.data_store.reset!
   end
 end
 
@@ -28,8 +28,6 @@ def test_repository(&block)
   Class.new do
     include Curator::Repository
 
-    self.data_store = Curator::Riak::Test::ResettableDataStore
-
     def self.name
       "TestModelRepository"
     end