Linking billable_periods to the relevant account_invoice upon creation

app/jobs/update_billable_periods.rb

@@ -71,9 +71,10 @@ class UpdateBillablePeriods
     owner_id = enterprise.owner_id
     sells = enterprise.sells
     orders = Spree::Order.where('distributor_id = (?) AND completed_at >= (?) AND completed_at < (?)', enterprise.id, begins_at, ends_at)
+    account_invoice = AccountInvoice.find_or_create_by_user_id_and_year_and_month(owner_id, begins_at.year, begins_at.month)
 
-    billable_period = BillablePeriod.where(begins_at: begins_at, enterprise_id: enterprise.id).first
-    billable_period ||= BillablePeriod.new(begins_at: begins_at, enterprise_id: enterprise.id)
+    billable_period = BillablePeriod.where(account_invoice_id: account_invoice.id, begins_at: begins_at, enterprise_id: enterprise.id).first
+    billable_period ||= BillablePeriod.new(account_invoice_id: account_invoice.id, begins_at: begins_at, enterprise_id: enterprise.id)
     billable_period.update_attributes({
       ends_at: ends_at,
       sells: sells,

app/models/account_invoice.rb

@@ -1,5 +1,6 @@
 class AccountInvoice < ActiveRecord::Base
-  belongs_to :user
-  belongs_to :order
-  attr_accessible :issued_at, :month, :year
+  belongs_to :user, class_name: "Spree::User"
+  belongs_to :order, class_name: "Spree::Order"
+  attr_accessible :user_id, :order_id, :issued_at, :month, :year
+  has_many :billable_periods
 end

app/models/billable_period.rb

@@ -1,7 +1,8 @@
 class BillablePeriod < ActiveRecord::Base
   belongs_to :enterprise
+  belongs_to :owner, class_name: 'Spree::User'
+  belongs_to :account_invoice
   has_one :adjustment, :as => :source, class_name: "Spree::Adjustment" #, :dependent => :destroy
-  belongs_to :owner, class_name: 'Spree::User', foreign_key: :owner_id
 
   default_scope where(deleted_at: nil)
 

spec/factories.rb

@@ -222,8 +222,18 @@ FactoryGirl.define do
     sells { 'any' }
     trial { false }
     enterprise
-    owner { FactoryGirl.create :user }
+    owner { enterprise.owner }
     turnover { rand(100000).to_f/100 }
+    account_invoice do
+      AccountInvoice.where(user_id: owner_id, year: begins_at.year, month: begins_at.month).first ||
+      FactoryGirl.create(user: owner, year: begins_at.year, month: begins_at.month)
+    end
+  end
+
+  factory :account_invoice do
+    user { FactoryGirl.create :user }
+    year { 2000 + rand(100) }
+    month { 1 + rand(12) }
   end
 end
 

spec/jobs/update_billable_periods_spec.rb

@@ -7,6 +7,7 @@ end
 describe UpdateBillablePeriods do
   describe "unit specs" do
     let!(:start_of_july) { Time.now.beginning_of_year + 6.months }
+    let!(:year) { Time.now.year }
 
     let!(:updater) { UpdateBillablePeriods.new }
 
@@ -44,7 +45,7 @@ describe UpdateBillablePeriods do
       end
 
       context "when a specfic year and month are passed as arguments" do
-        let!(:updater) { UpdateBillablePeriods.new(Time.now.year, 6) }
+        let!(:updater) { UpdateBillablePeriods.new(year, 6) }
 
         before do
           allow(updater).to receive(:split_for_trial)
@@ -494,12 +495,17 @@ describe UpdateBillablePeriods do
     # Chose july to test with because June has 30 days and so is easy to calculate end date for shop trial
     let!(:start_of_july) { Time.now.beginning_of_year + 6.months }
 
+    let!(:year) { Time.now.year }
+
     let!(:enterprise) { create(:supplier_enterprise, sells: 'any') }
 
     let!(:original_owner) { enterprise.owner }
 
     let!(:new_owner) { create(:user) }
 
+    let!(:account_invoice1) { create(:account_invoice, user: original_owner, year: year, month: 7)}
+    let!(:account_invoice2) { create(:account_invoice, user: new_owner, year: year, month: 7)}
+
     # This BP was updated before the current run and so should be marked for deletion at the end of the run
     let!(:obsolete_bp) { create(:billable_period, enterprise: enterprise, updated_at: start_of_july + 10.days, begins_at: start_of_july + 6.5.days, ends_at: start_of_july + 10.days ) }
 
@@ -568,15 +574,20 @@ describe UpdateBillablePeriods do
         expect(obsolete_bp.reload.deleted_at).to_not be_nil
 
         bp_to_overwrite.reload
+
         expect(bp_to_overwrite.sells).to eq 'own'
         expect(bp_to_overwrite.trial).to be true
         expect(bp_to_overwrite.owner).to eq original_owner
         expect(bp_to_overwrite.begins_at).to eq start_of_july + 10.days
         expect(bp_to_overwrite.ends_at).to eq start_of_july + 12.days
         expect(bp_to_overwrite.turnover).to eq order6.total
+        expect(bp_to_overwrite.account_invoice).to eq account_invoice1
 
         expect(billable_periods.count).to eq 9
 
+        expect(account_invoice1.billable_periods.sort).to eq billable_periods.sort.select{ |bp| bp.owner == original_owner }
+        expect(account_invoice2.billable_periods.sort).to eq billable_periods.sort.select{ |bp| bp.owner == new_owner }
+
         expect(billable_periods.map(&:begins_at)).to eq [
           start_of_july + 2.days,
           start_of_july + 4.days,