Clean view context when controller changes

lib/draper/view_context.rb

@@ -20,8 +20,10 @@ def self.controller
       RequestStore.store[:current_controller]
     end
 
-    # Sets the current controller.
+    # Sets the current controller. Clears view context when we are setting
+    # different controller.
     def self.controller=(controller)
+      clear! if RequestStore.store[:current_controller] != controller
       RequestStore.store[:current_controller] = controller
     end
 

spec/draper/view_context_spec.rb

@@ -32,6 +32,34 @@ module Draper
         ViewContext.controller = :stored_controller
         expect(store[:current_controller]).to be :stored_controller
       end
+
+      it "cleans context when controller changes" do
+        store = {
+          current_controller: :stored_controller,
+          current_view_context: :stored_view_context
+        }
+
+        allow(RequestStore).to receive_messages store: store
+
+        ViewContext.controller = :other_stored_controller
+
+        expect(store).to include(current_controller: :other_stored_controller)
+        expect(store).not_to include(:current_view_context)
+      end
+
+      it "doesn't clean context when controller is the same" do
+        store = {
+          current_controller: :stored_controller,
+          current_view_context: :stored_view_context
+        }
+
+        allow(RequestStore).to receive_messages store: store
+
+        ViewContext.controller = :stored_controller
+
+        expect(store).to include(current_controller: :stored_controller)
+        expect(store).to include(current_view_context: :stored_view_context)
+      end
     end
 
     describe ".current" do