Associate Locations created via import

Add the model associations and populate the field when importing.

app/models/immunisation_import.rb

@@ -27,6 +27,9 @@ class ImmunisationImport < ApplicationRecord
   has_many :vaccination_records,
            dependent: :restrict_with_exception,
            foreign_key: :imported_from_id
+  has_many :locations,
+           dependent: :restrict_with_exception,
+           foreign_key: :imported_from_id
 
   EXPECTED_HEADERS = %w[
     ANATOMICAL_SITE
@@ -73,8 +76,10 @@ def process!(patient_session:)
         .map(&:to_location)
         .uniq(&:urn)
         .reject(&:invalid?)
-        .map(&:attributes)
-        .tap { Location.create! _1 }
+        .each do |location|
+          location.imported_from = self
+          location.save!
+        end
 
       rows
         .map(&:to_vaccination_record)

app/models/location.rb

@@ -32,6 +32,7 @@ class Location < ApplicationRecord
   has_many :sessions
   has_many :patients
   has_many :consent_forms, through: :sessions
+  belongs_to :imported_from, class_name: "ImmunisationImport", optional: true
 
   validates :name, presence: true
   validates :urn, presence: true, uniqueness: true

spec/models/immunisation_import_spec.rb

@@ -86,7 +86,7 @@
       expect { immunisation_import.process!(patient_session:) }.to change(
         immunisation_import.vaccination_records,
         :count
-      ).by(11).and change(Location, :count).by(4)
+      ).by(11).and change(immunisation_import.locations, :count).by(4)
     end
   end
 end