4 files changed, 38 insertions, 21 deletions

diff --git a/app/models/chouette/journey_pattern.rb b/app/models/chouette/journey_pattern.rb
index c57d42e71..34b0d9345 100644
--- a/app/models/chouette/journey_pattern.rb
+++ b/app/models/chouette/journey_pattern.rb
@@ -25,9 +25,7 @@ class Chouette::JourneyPattern < Chouette::TridentActiveRecord
       state.each do |item|
         item.delete('errors')
         jp = find_by(objectid: item['object_id']) || state_create_instance(route, item)
-        if item['deletable'] && jp.persisted?
-          next if jp.destroy
-        end
+        next if item['deletable'] && jp.persisted? && jp.destroy
 
         # Update attributes and stop_points associations
         jp.update_attributes(state_permited_attributes(item))
@@ -40,7 +38,7 @@ class Chouette::JourneyPattern < Chouette::TridentActiveRecord
         raise ActiveRecord::Rollback
       end
     end
-    # clean
+
     state.map {|item| item.delete('new_record')}
     state.delete_if {|item| item['deletable']}
   end
@@ -54,8 +52,8 @@ class Chouette::JourneyPattern < Chouette::TridentActiveRecord
   end
 
   def self.state_create_instance route, item
-    jp = route.journey_patterns.create(state_permited_attributes(item))
     # Flag new record, so we can unset object_id if transaction rollback
+    jp = route.journey_patterns.create(state_permited_attributes(item))
     item['object_id']  = jp.objectid
     item['new_record'] = true
     jp
diff --git a/app/models/chouette/vehicle_journey.rb b/app/models/chouette/vehicle_journey.rb
index f22ddd390..0aab9d0ed 100644
--- a/app/models/chouette/vehicle_journey.rb
+++ b/app/models/chouette/vehicle_journey.rb
@@ -74,38 +74,41 @@ module Chouette
       end
     end
 
-    def update_company_from_state state
-      if state['company']['id'] != state['company_id']
-        self.company = Company.find(state['company']['id'])
-        state['company_id'] = self.company.id
-        self.save
-      end
-    end
-
     def self.state_update route, state
       transaction do
         state.each do |item|
           item.delete('errors')
-          vj = find_by(objectid: item['objectid'])
+          vj = find_by(objectid: item['objectid']) || state_create_instance(route, item)
           next if item['deletable'] && vj.persisted? && vj.destroy
 
           vj.update_vjas_from_state(item['vehicle_journey_at_stops'])
           vj.update_attributes(state_permited_attributes(item))
-          vj.update_company_from_state(item) if item['company']
-
           item['errors'] = vj.errors if vj.errors.any?
         end
         if state.any? {|item| item['errors']}
+          state.map {|item| item.delete('objectid') if item['new_record']}
           raise ::ActiveRecord::Rollback
         end
       end
 
-      # Cleanup
+      state.map {|item| item.delete('new_record')}
       state.delete_if {|item| item['deletable']}
     end
 
+    def self.state_create_instance route, item
+      # Flag new record, so we can unset object_id if transaction rollback
+      vj = route.vehicle_journeys.create(state_permited_attributes(item))
+      item['objectid']   = vj.objectid
+      item['new_record'] = true
+      vj
+    end
+
     def self.state_permited_attributes item
-      item.slice('published_journey_identifier', 'published_journey_name').to_hash
+      attrs = item.slice('published_journey_identifier', 'published_journey_name', 'journey_pattern_id', 'company_id').to_hash
+      ['company', 'journey_pattern'].map do |association|
+        attrs["#{association}_id"] = item[association]['id'] if item[association]
+      end
+      attrs
     end
 
     def increasing_times
diff --git a/app/views/vehicle_journeys/show.rabl b/app/views/vehicle_journeys/show.rabl
index 398bb3b16..86edfafa8 100644
--- a/app/views/vehicle_journeys/show.rabl
+++ b/app/views/vehicle_journeys/show.rabl
@@ -1,6 +1,6 @@
 object @vehicle_journey
 
-[:objectid, :published_journey_name, :published_journey_identifier].each do |attr|
+[:objectid, :published_journey_name, :published_journey_identifier, :company_id].each do |attr|
   attributes attr, :unless => lambda { |m| m.send( attr).nil?}
 end
 
diff --git a/spec/models/chouette/vehicle_journey_spec.rb b/spec/models/chouette/vehicle_journey_spec.rb
index cb43549d0..017aa964b 100644
--- a/spec/models/chouette/vehicle_journey_spec.rb
+++ b/spec/models/chouette/vehicle_journey_spec.rb
@@ -28,11 +28,27 @@ describe Chouette::VehicleJourney, :type => :model do
     let(:state)           { vehicle_journey_to_state(vehicle_journey) }
     let(:collection)      { [state] }
 
+    it 'should create new vj from state' do
+      new_vj = build(:vehicle_journey, objectid: nil, published_journey_name: 'dummy', route: route, journey_pattern: journey_pattern)
+      collection << vehicle_journey_to_state(new_vj)
+      expect {
+        Chouette::VehicleJourney.state_update(route, collection)
+      }.to change {Chouette::VehicleJourney.count}.by(1)
+      expect(collection.last['objectid']).not_to be_nil
+
+      vj = Chouette::VehicleJourney.find_by(objectid: collection.last['objectid'])
+      expect(vj.published_journey_name).to eq 'dummy'
+    end
+
+    it 'should update vj journey_pattern' do
+      state['journey_pattern'] = create(:journey_pattern).attributes.slice('id', 'name', 'objectid')
+      Chouette::VehicleJourney.state_update(route, collection)
+      expect(vehicle_journey.reload.journey_pattern_id).to eq state['journey_pattern']['id']
+    end
+
     it 'should update vj company' do
       state['company'] = create(:company).attributes.slice('id', 'name', 'objectid')
       Chouette::VehicleJourney.state_update(route, collection)
-
-      expect(state['company_id']).to eq state['company']['id']
       expect(vehicle_journey.reload.company_id).to eq state['company']['id']
     end