Make DataFrabric.activate() nesting-friendly.

lib/data_fabric.rb

@@ -48,15 +48,26 @@ def self.init
 
   def self.activate_shard(shards, &block)
     ensure_setup
-    shards.each do |key, value|
+    old_shard_values = {}
+    newly_added_shard_keys = []
+    shards.each_pair do |key, value|
+      key = key.to_s
+      if Thread.current[:shards].has_key?(key)
+        old_shard_values[key] = Thread.current[:shards][key]
+      else
+        newly_added_shard_keys << key
+      end
       Thread.current[:shards][key.to_s] = value.to_s
     end
     if block_given?
       begin
         yield
       ensure
-        shards.each do |key, value|
-          Thread.current[:shards].delete(key.to_s)
+        old_shard_values.each_pair do |key, value|
+          Thread.current[:shards][key] = value
+        end
+        newly_added_shard_keys.each do |key|
+          Thread.current[:shards].delete(key)
         end
       end
     end

test/shard_test.rb

@@ -21,4 +21,31 @@ def test_activation_in_one_thread_does_not_change_another
       assert_equal 'dallas', DataFabric.active_shard(:city)
     end.join
   end
-end
+
+  def test_activations_can_be_nested
+    DataFabric.activate_shard(:name => 'Belldandy') do
+      DataFabric.activate_shard(:technique => 'Thousand Years of Pain') do
+        DataFabric.activate_shard(:name => 'Lelouche') do
+          DataFabric.activate_shard(:technique => 'Shadow Replication') do
+            assert_equal 'Shadow Replication', DataFabric.active_shard(:technique)
+            assert_equal 'Lelouche', DataFabric.active_shard(:name)
+          end
+          assert_equal 'Thousand Years of Pain', DataFabric.active_shard(:technique)
+          assert_equal 'Lelouche', DataFabric.active_shard(:name)
+        end
+        assert_equal 'Thousand Years of Pain', DataFabric.active_shard(:technique)
+        assert_equal 'Belldandy', DataFabric.active_shard(:name)
+      end
+      assert_raises ArgumentError do
+        DataFabric.active_shard(:technique)
+      end
+      assert_equal 'Belldandy', DataFabric.active_shard(:name)
+    end
+    assert_raises ArgumentError do
+      DataFabric.active_shard(:technique)
+    end
+    assert_raises ArgumentError do
+      DataFabric.active_shard(:name)
+    end
+  end
+end