makes ballot's back link point to referer

app/controllers/budgets/ballot/lines_controller.rb

@@ -7,8 +7,8 @@ module Budgets
       before_action :load_ballot
       before_action :load_tag_cloud
       before_action :load_categories
-
       before_action :load_investments
+      before_action :load_ballot_referer
 
       load_and_authorize_resource :budget
       load_and_authorize_resource :ballot, class: "Budget::Ballot", through: :budget
@@ -73,6 +73,10 @@ module Budgets
           @categories = ActsAsTaggableOn::Tag.where("kind = 'category'").order(:name)
         end
 
+        def load_ballot_referer
+          @ballot_referer = session[:ballot_referer]
+        end
+
     end
   end
 end

app/controllers/budgets/ballots_controller.rb

@@ -3,9 +3,11 @@ module Budgets
     before_action :authenticate_user!
     load_and_authorize_resource :budget
     before_action :load_ballot
+    after_action :store_referer, only: [:show]
 
     def show
       authorize! :show, @ballot
+      session[:ballot_referer] = request.referer
       render template: "budgets/ballot/show"
     end
 
@@ -16,5 +18,9 @@ module Budgets
         @ballot = @budget.balloting? ? query.first_or_create : query.first_or_initialize
       end
 
+      def store_referer
+        session[:ballot_referer] = request.referer
+      end
+
   end
 end

app/views/budgets/ballot/_ballot.html.erb

@@ -1,6 +1,6 @@
 <div class="expanded budget no-margin-top padding">
   <div class="row">
-    <%= render 'shared/back_link' %>
+    <%= back_link_to @ballot_referer %>
 
     <h1 class="text-center"><%= t("budgets.ballots.show.title") %></h1>
 

spec/features/budgets/ballots_spec.rb

@@ -350,7 +350,7 @@ feature 'Ballots' do
 
   end
 
-  scenario 'Removing spending proposals from ballot', :js do
+  scenario 'Removing investments from ballot', :js do
     investment = create(:budget_investment, :selected, price: 10, heading: new_york)
     ballot = create(:budget_ballot, user: user, budget: budget)
     ballot.investments << investment
@@ -368,7 +368,7 @@ feature 'Ballots' do
     expect(page).to have_content("You have voted 0 proposals")
   end
 
-  scenario 'Removing spending proposals from ballot (sidebar)', :js do
+  scenario 'Removing investments from ballot (sidebar)', :js do
     investment1 = create(:budget_investment, :selected, price: 10000, heading: new_york)
     investment2 = create(:budget_investment, :selected, price: 20000, heading: new_york)
 
@@ -405,6 +405,28 @@ feature 'Ballots' do
     end
   end
 
+  scenario 'Back link after removing an investment from Ballot', :js do
+    investment = create(:budget_investment, :selected, heading: new_york, price: 10)
+
+    login_as(user)
+    visit budget_investments_path(budget, heading_id: new_york.id)
+    add_to_ballot(investment)
+
+    click_link "Check my ballot"
+
+    expect(page).to have_content("You have voted one proposal")
+
+    within("#budget_investment_#{investment.id}") do
+      find(".remove-investment-project").trigger('click')
+    end
+
+    expect(page).to have_content("You have voted 0 proposals")
+
+    click_link "Go back"
+
+    expect(page).to have_current_path(budget_investments_path(budget, heading_id: new_york.id))
+  end
+
   context 'Permissions' do
 
     scenario 'User not logged in', :js do