From 7ebff4e9b5e6890c0a7e4260a90fe0db66656e43 Mon Sep 17 00:00:00 2001
From: Austen Ito <austen.ito@gmail.com>
Date: Fri, 3 Mar 2017 19:39:27 -0500
Subject: [PATCH] Fix adding custom context to current context (#54)

* Fix adding custom context to current context

* The data hash would lose custom context information
  since merging was not happening in place.

* Add CurrentContext specs
---
 lib/timber/current_context.rb       | 61 ++++++++++++++++++-----------
 spec/timber/current_context_spec.rb | 57 +++++++++++++++++++++++++++
 2 files changed, 96 insertions(+), 22 deletions(-)
 create mode 100644 spec/timber/current_context_spec.rb

diff --git a/lib/timber/current_context.rb b/lib/timber/current_context.rb
index aee0903c..aa36e6a8 100644
--- a/lib/timber/current_context.rb
+++ b/lib/timber/current_context.rb
@@ -23,13 +23,17 @@ def add(*args)
         instance.add(*args)
       end
 
+      def hash(*args)
+        instance.hash(*args)
+      end
+
       # Convenience method for {#remove}. See {#remove} for full details and examples.
       def remove(*args)
         instance.remove(*args)
       end
 
-      def hash(*args)
-        instance.hash(*args)
+      def reset(*args)
+        instance.reset(*args)
       end
     end
 
@@ -56,18 +60,11 @@ def hash(*args)
     #
     # @example Adding multiple contexts
     #   Timber::CurrentContext.with(context1, context2) { ... }
-    def with(*contexts)
-      add(*contexts)
+    def with(*objects)
+      add(*objects)
       yield
     ensure
-      contexts.each do |context|
-        if context.keyspace == :custom
-          # Custom contexts are merged and should be removed the same
-          hash[context.keyspace].delete(context.type)
-        else
-          remove(context)
-        end
-      end
+      remove(*objects)
     end
 
     # Adds contexts but does not remove them. See {#with} for automatic maintenance and {#remove}
@@ -83,31 +80,51 @@ def add(*objects)
         if key == :custom
           # Custom contexts are merged into the space
           hash[key] ||= {}
-          hash[key].merge(json)
+          hash[key].merge!(json)
         else
           hash[key] = json
         end
       end
     end
 
-    # Removes a context. This must be a {Timber::Context} type. See {Timber::Contexts} for a list.
-    # If you wish to remove by key, or some other way, use {#hash} and modify the hash accordingly.
-    def remove(*contexts)
-      contexts.each do |context|
-        hash.delete(context.keyspace)
-      end
-    end
-
     # The internal hash that is maintained. It is recommended that you use {#with} and {#add}
     # for hash maintenance.
     def hash
       Thread.current[THREAD_NAMESPACE] ||= {}
     end
 
+    # Removes a context. If you wish to remove by key, or some other way, use {#hash} and
+    # modify the hash accordingly.
+    def remove(*objects)
+      objects.each do |object|
+        if object.is_a?(Symbol)
+          hash.delete(object)
+        else
+          context = Contexts.build(object)
+
+          if context.keyspace == :custom
+            # Custom contexts are merged and should be removed the same
+            hash[context.keyspace].delete(context.type)
+            if hash[context.keyspace] == {}
+              hash.delete(context.keyspace)
+            end
+          else
+            hash.delete(context.keyspace)
+          end
+        end
+      end
+    end
+
+    # Resets the context to be blank. Use this carefully! This will remove *any* context,
+    # include context that is automatically included with Timber.
+    def reset
+      hash.clear
+    end
+
     # Snapshots the current context so that you get a moment in time representation of the context,
     # since the context can change as execution proceeds.
     def snapshot
       hash.clone
     end
   end
-end
\ No newline at end of file
+end
diff --git a/spec/timber/current_context_spec.rb b/spec/timber/current_context_spec.rb
new file mode 100644
index 00000000..cf38b7aa
--- /dev/null
+++ b/spec/timber/current_context_spec.rb
@@ -0,0 +1,57 @@
+require "spec_helper"
+
+describe Timber::CurrentContext, :rails_23 => true do
+  describe ".add" do
+    after(:each) do
+      described_class.reset
+    end
+
+    it "should add the context" do
+      expect(described_class.hash).to eq({})
+
+      described_class.add({build: {version: "1.0.0"}})
+      expect(described_class.hash).to eq({:custom=>{:build=>{:version=>"1.0.0"}}})
+
+      described_class.add({testing: {key: "value"}})
+      expect(described_class.hash).to eq({:custom=>{:build=>{:version=>"1.0.0"}, :testing=>{:key=>"value"}}})
+    end
+  end
+
+  describe ".remove" do
+    it "should remove the context by object" do
+      context = {:build=>{:version=>"1.0.0"}}
+      described_class.add(context)
+      expect(described_class.hash).to eq({:custom => context})
+
+      described_class.remove(context)
+      expect(described_class.hash).to eq({})
+    end
+
+    it "should remove context by key" do
+      context = {:build=>{:version=>"1.0.0"}}
+      described_class.add(context)
+      expect(described_class.hash).to eq({:custom => context})
+
+      described_class.remove(:custom)
+      expect(described_class.hash).to eq({})
+    end
+  end
+
+  describe ".with" do
+    it "should merge the context and cleanup on block exit" do
+      expect(described_class.hash).to eq({})
+
+      described_class.with({build: {version: "1.0.0"}}) do
+        expect(described_class.hash).to eq({:custom=>{:build=>{:version=>"1.0.0"}}})
+
+        described_class.with({testing: {key: "value"}}) do
+          expect(described_class.hash).to eq({:custom=>{:build=>{:version=>"1.0.0"}, :testing=>{:key=>"value"}}})
+        end
+
+        expect(described_class.hash).to eq({:custom=>{:build=>{:version=>"1.0.0"}}})
+      end
+
+      expect(described_class.hash).to eq({})
+    end
+  end
+end
\ No newline at end of file