Simplify budget stats data generation

Configuring associations in the factories makes it easier to generate
headings, polls and ballots in the tests.

spec/factories/budgets.rb

@@ -74,13 +74,15 @@ FactoryBot.define do
   end
 
   factory :budget_heading, class: "Budget::Heading" do
-    association :group, factory: :budget_group
     sequence(:name) { |n| "Heading #{n}" }
     price 1000000
     population 1234
     latitude "40.416775"
     longitude "-3.703790"
 
+    transient { budget nil }
+    group { association :budget_group, budget: budget || association(:budget) }
+
     trait :drafting_budget do
       association :group, factory: [:budget_group, :drafting_budget]
     end
@@ -88,7 +90,7 @@ FactoryBot.define do
 
   factory :budget_investment, class: "Budget::Investment" do
     sequence(:title) { |n| "Budget Investment #{n} title" }
-    association :heading, factory: :budget_heading
+    heading { association :budget_heading, budget: budget }
     association :author, factory: :user
     description          "Spend money on this"
     price                10
@@ -189,8 +191,13 @@ FactoryBot.define do
   end
 
   factory :budget_ballot_line, class: "Budget::Ballot::Line" do
-    association :ballot, factory: :budget_ballot
     association :investment, factory: :budget_investment
+
+    transient { user nil }
+
+    ballot do
+      association :budget_ballot, budget: investment.budget, user: user || association(:user)
+    end
   end
 
   factory :budget_reclassified_vote, class: "Budget::ReclassifiedVote" do

spec/factories/polls.rb

@@ -88,10 +88,13 @@ FactoryBot.define do
   end
 
   factory :poll_voter, class: "Poll::Voter" do
-    poll
     association :user, :level_two
     from_web
 
+    transient { budget nil }
+
+    poll { budget&.poll || association(:poll, budget: budget) }
+
     trait :from_web do
       origin "web"
       token SecureRandom.hex(32)

spec/models/budget/stats_spec.rb

@@ -4,8 +4,7 @@ describe Budget::Stats do
 
   before do
     @budget = create(:budget)
-    @group = create(:budget_group, budget: @budget)
+    @heading = create(:budget_heading, budget: @budget, price: 1000)
-    @heading = create(:budget_heading, group: @group, price: 1000)
 
     @investment1 = create(:budget_investment, :selected, author: create(:user, gender: "female"),
                           heading: @heading, price: 200, ballot_lines_count: 900, winner: true)
@@ -22,20 +21,15 @@ describe Budget::Stats do
     @support1 = create(:vote, votable: @investment1, voter: create(:user, gender: "male"))
     @support2 = create(:vote, votable: @investment2, voter: create(:user))
 
-    @budget_ballot1 = create(:budget_ballot, budget: @budget, user: create(:user, gender: "female",
+    @budget_ballot_line1 = create(:budget_ballot_line,
-                             date_of_birth: 54.years.ago))
+                                  user: create(:user, gender: "female", date_of_birth: 54.years.ago),
-    @budget_ballot2 = create(:budget_ballot, budget: @budget, user: create(:user, gender: "female"))
-    @budget_ballot3 = create(:budget_ballot, budget: @budget, user: create(:user, gender: "male"))
-
-    @budget_ballot_line1 = create(:budget_ballot_line, ballot: @budget_ballot1,
                                   investment: @investment1)
-    @budget_ballot_line2 = create(:budget_ballot_line, ballot: @budget_ballot2,
+    @budget_ballot_line2 = create(:budget_ballot_line, user: create(:user, gender: "female"),
                                   investment: @investment2)
-    @budget_ballot_line3 = create(:budget_ballot_line, ballot: @budget_ballot3,
+    @budget_ballot_line3 = create(:budget_ballot_line, user: create(:user, gender: "male"),
                                   investment: @investment3)
 
-    @poll = create(:poll, budget: @budget)
+    @poll_voter = create(:poll_voter, :from_booth, budget: @budget)
-    @poll_voter = create(:poll_voter, :from_booth, poll: @poll)
 
     @budget_ballot4 = create(:budget_ballot, budget: @budget, physical: true, user: nil)
     @budget_ballot_line4 = create(:budget_ballot_line, ballot: @budget_ballot4,