Adds more controller variables and uses them in investment controller

app/controllers/management/budgets/investments_controller.rb

@@ -6,6 +6,8 @@ class Management::Budgets::InvestmentsController < Management::BaseController
 
   def index
     @investments = apply_filters_and_search(Budget::Investment).order(cached_votes_up: :desc).page(params[:page]).for_render
+    @investment_ids = @investments.pluck(:id)
+    set_investment_ballots(@investments)
     set_budget_investment_votes(@investments)
   end
 
@@ -26,11 +28,13 @@ class Management::Budgets::InvestmentsController < Management::BaseController
   end
 
   def show
+    set_investment_ballots(@investment)
     set_budget_investment_votes(@investment)
   end
 
   def vote
     @investment.register_vote(managed_user, 'yes')
+    set_investment_ballots(@investment)
     set_budget_investment_votes(@investment)
   end
 
@@ -42,6 +46,13 @@ class Management::Budgets::InvestmentsController < Management::BaseController
 
   private
 
+    def set_investment_ballots(investments)
+      @investment_ballots = {}
+      Budget.where(id: Array.wrap(investments).map(&:budget_id).uniq).each do |budget|
+        @investment_ballots[budget] = Budget::Ballot.where(user: current_user, budget: budget).first_or_create
+      end
+    end
+
     def set_budget_investment
       @investment = Budget::Investment.find(params[:id])
     end

app/views/management/budgets/investments/index.html.erb

@@ -18,7 +18,12 @@
       </div>
 
       <% @investments.each do |investment| %>
-        <%= render partial: '/budgets/investments/investment', locals: { investment: investment, investment_votes: @investment_votes } %>
+        <%= render partial: '/budgets/investments/investment', locals: {
+          investment: investment,
+          investment_ids: @investment_ids,
+          investment_votes: @investment_votes,
+          investment_ballots: @investment_ballots
+        } %>
       <% end %>
 
       <%= paginate @investments %>