diff --git a/app/models/budget/ballot.rb b/app/models/budget/ballot.rb index 17e514d2c..6a4047aff 100644 --- a/app/models/budget/ballot.rb +++ b/app/models/budget/ballot.rb @@ -66,6 +66,7 @@ class Budget end def heading_for_group(group) + return nil unless has_lines_in_group?(group) self.investments.where(group: group).first.heading end diff --git a/spec/models/budget/ballot_spec.rb b/spec/models/budget/ballot_spec.rb index 9d4ff9608..3a9f9f232 100644 --- a/spec/models/budget/ballot_spec.rb +++ b/spec/models/budget/ballot_spec.rb @@ -79,7 +79,7 @@ describe Budget::Ballot do describe "#heading_for_group" do - it "returns the heading with investments for a group" do + it "returns the heading with balloted investments for a group" do budget = create(:budget) group = create(:budget_group, budget: budget) @@ -95,6 +95,18 @@ describe Budget::Ballot do 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