From e2b536e8f8a7d12e5de5a52721aa212bcade9826 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Javi=20Mart=C3=ADn?= <javim@elretirao.net>
Date: Sat, 28 Sep 2019 20:25:32 +0200
Subject: [PATCH] Simplify adding investments to ballots in specs

---
 spec/factories/budgets.rb             | 14 ++++++++++++++
 spec/models/budget/ballot_spec.rb     | 15 +++++----------
 spec/models/budget/investment_spec.rb |  6 ++----
 3 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/spec/factories/budgets.rb b/spec/factories/budgets.rb
index 786ee9922..3f6e637ba 100644
--- a/spec/factories/budgets.rb
+++ b/spec/factories/budgets.rb
@@ -194,12 +194,18 @@ FactoryBot.define do
     transient do
       voters { [] }
       followers { [] }
+      ballots { [] }
       balloters { [] }
     end
 
     after(:create) do |investment, evaluator|
       evaluator.voters.each { |voter| create(:vote, votable: investment, voter: voter) }
       evaluator.followers.each { |follower| create(:follow, followable: investment, user: follower) }
+
+      evaluator.ballots.each do |ballot|
+        create(:budget_ballot_line, investment: investment, ballot: ballot)
+      end
+
       evaluator.balloters.each do |balloter|
         create(:budget_ballot_line, investment: investment, user: balloter)
       end
@@ -219,6 +225,14 @@ FactoryBot.define do
   factory :budget_ballot, class: "Budget::Ballot" do
     association :user, factory: :user
     budget
+
+    transient { investments { [] } }
+
+    after(:create) do |ballot, evaluator|
+      evaluator.investments.each do |investment|
+        create(:budget_ballot_line, investment: investment, ballot: ballot)
+      end
+    end
   end
 
   factory :budget_ballot_line, class: "Budget::Ballot::Line" do
diff --git a/spec/models/budget/ballot_spec.rb b/spec/models/budget/ballot_spec.rb
index 039706816..bcfcc6c28 100644
--- a/spec/models/budget/ballot_spec.rb
+++ b/spec/models/budget/ballot_spec.rb
@@ -13,11 +13,8 @@ describe Budget::Ballot do
 
     it "is not valid with the same investment twice" do
       budget = create(:budget)
-      heading = create(:budget_heading, budget: budget)
-      investment = create(:budget_investment, :selected, heading: heading)
-
-      ballot = create(:budget_ballot, budget: budget)
-      ballot.investments << investment
+      investment = create(:budget_investment, :selected, budget: budget)
+      ballot = create(:budget_ballot, budget: budget.reload, investments: [investment])
 
       expect { ballot.investments << investment }.to raise_error(ActiveRecord::RecordNotUnique)
     end
@@ -103,16 +100,14 @@ describe Budget::Ballot do
 
     it "returns the heading with balloted investments for a group" do
       budget = create(:budget)
+      ballot = create(:budget_ballot, budget: 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
+      create(:budget_investment, :selected, heading: heading1)
+      create(:budget_investment, :selected, heading: heading2, ballots: [ballot])
 
       expect(ballot.heading_for_group(group)).to eq heading2
     end
diff --git a/spec/models/budget/investment_spec.rb b/spec/models/budget/investment_spec.rb
index aa7ae9845..94ab490d8 100644
--- a/spec/models/budget/investment_spec.rb
+++ b/spec/models/budget/investment_spec.rb
@@ -976,8 +976,7 @@ describe Budget::Investment do
 
           inv1 = create(:budget_investment, :selected, budget: budget, heading: california)
           inv2 = create(:budget_investment, :selected, budget: budget, heading: new_york)
-          ballot = create(:budget_ballot, user: user, budget: budget)
-          ballot.investments << inv1
+          ballot = create(:budget_ballot, user: user, budget: budget, investments: [inv1])
 
           expect(inv2.reason_for_not_being_ballotable_by(user, ballot)).to eq(:different_heading_assigned_html)
         end
@@ -989,8 +988,7 @@ describe Budget::Investment do
           inv1 = create(:budget_investment, :selected, budget: budget, heading: carabanchel, price: 30)
           inv2 = create(:budget_investment, :selected, budget: budget, heading: carabanchel, price: 10)
 
-          ballot = create(:budget_ballot, user: user, budget: budget)
-          ballot.investments << inv1
+          ballot = create(:budget_ballot, user: user, budget: budget, investments: [inv1])
 
           expect(inv2.reason_for_not_being_ballotable_by(user, ballot)).to eq(:not_enough_money_html)
         end