Merge pull request #1512 from consul/budgets-review-my-votes

Displays the correct amount spent in review my ballot

app/models/budget/ballot.rb

@@ -66,7 +66,8 @@ class Budget
     end
 
     def heading_for_group(group)
-      self.headings.where(group: group).first
+      return nil unless has_lines_in_group?(group)
+      self.investments.where(group: group).first.heading
     end
 
   end

spec/models/budget/ballot_spec.rb

@@ -5,10 +5,13 @@ describe Budget::Ballot do
   describe "#amount_spent" do
     it "returns the total amount spent in investments" do
       budget = create(:budget)
+
       group1 = create(:budget_group, budget: budget)
       group2 = create(:budget_group, budget: budget)
+
       heading1 = create(:budget_heading, group: group1, price: 100000)
       heading2 = create(:budget_heading, group: group2, price: 200000)
+
       inv1 = create(:budget_investment, :selected, price: 10000, heading: heading1)
       inv2 = create(:budget_investment, :selected, price: 20000, heading: heading2)
 
@@ -24,10 +27,13 @@ describe Budget::Ballot do
 
     it "returns the amount spent on all investments assigned to a specific heading" do
       budget = create(:budget)
+
       group1 = create(:budget_group, budget: budget)
       group2 = create(:budget_group, budget: budget)
+
       heading1 = create(:budget_heading, group: group1, price: 100000)
       heading2 = create(:budget_heading, group: group2, price: 200000)
+
       inv1 = create(:budget_investment, :selected, price: 10000, heading: heading1)
       inv2 = create(:budget_investment, :selected, price: 20000, heading: heading2)
       inv3 = create(:budget_investment, :selected, price: 40000, heading: heading1)
@@ -48,10 +54,13 @@ describe Budget::Ballot do
   describe "#amount_available" do
     it "returns how much is left after taking some investments" do
       budget = create(:budget)
+
       group1 = create(:budget_group, budget: budget)
       group2 = create(:budget_group, budget: budget)
+
       heading1 = create(:budget_heading, group: group1, price: 1000)
       heading2 = create(:budget_heading, group: group2, price: 300)
+
       inv1 = create(:budget_investment, :selected, price: 100, heading: heading1)
       inv2 = create(:budget_investment, :selected, price: 200, heading: heading2)
       inv3 = create(:budget_investment, :selected, price: 400, heading: heading1)
@@ -68,4 +77,36 @@ describe Budget::Ballot do
     end
   end
 
+  describe "#heading_for_group" do
+
+    it "returns the heading with balloted investments for a group" do
+      budget = create(:budget)
+      group = create(:budget_group, budget: budget)
+
+      heading1 = create(:budget_heading, group: group)
+      heading2 = create(:budget_heading, group: group)
+
+      inv1 = create(:budget_investment, :selected, heading: heading1)
+      inv2 = create(:budget_investment, :selected, heading: heading2)
+
+      ballot = create(:budget_ballot, budget: budget)
+      ballot.investments << inv2
+
+      expect(ballot.heading_for_group(group)).to eq heading2
+    end
+
+    it "returns nil if there are no headings with balloted investments in a group" do
+      budget = create(:budget)
+      group = create(:budget_group, budget: budget)
+
+      heading1 = create(:budget_heading, group: group)
+      heading2 = create(:budget_heading, group: group)
+
+      ballot = create(:budget_ballot, budget: budget)
+
+      expect(ballot.heading_for_group(group)).to eq nil
+    end
+
+  end
+
 end