Updates to the data generators (#3478)

Incremental changes as a result of using these recently. See commits for
the changes.

Author: misaka

lib/generate/patient_imports.rb lib/generate/cohort_imports.rblib/generate/patient_imports.rb renamed to lib/generate/cohort_imports.rb

@@ -8,17 +8,41 @@
 #
 # Usage from the Rails console:
 #
-#     # By default it uses existing locations in the db.
-#     Generate::PatientImports.call
+# Create a cohort import of 1000 children for all the school sessions for the
+# org A9A5A in the local db:
+#
+#     Generate::CohortImports.call(patient_count: 1000)
+#
+# You can also generate a cohort import for sessions not in the local db.
+#
+#     Generate::CohortImports.call(
+#       patient_count: 1000,
+#       urns: ["123456", "987654"],
+#       school_year_groups: {
+#         "123456" => [-2, -1, 0, 1, 2, 3, 4, 5, 6],
+#         "987654" => [9, 10, 11, 12, 13]
+#       }
+#     )
+#
+# You can pull out the year groups with the following:
+#
+#     org = Organisation.find_by(ods_code: "A9A5A")
+#     org.locations.school.pluck(:urn, :year_groups) .to_h
 #
 module Generate
-  class PatientImports
-    attr_reader :ods_code, :organisation, :programme, :urns, :patient_count
+  class CohortImports
+    attr_reader :ods_code,
+                :organisation,
+                :programme,
+                :urns,
+                :patient_count,
+                :school_year_groups
 
     def initialize(
       ods_code: "A9A5A",
       programme: "hpv",
       urns: nil,
+      school_year_groups: nil,
       patient_count: 10
     )
       @organisation = Organisation.find_by(ods_code:)
@@ -30,19 +54,19 @@ def initialize(
             .select { it.urn.present? }
             .sample(3)
             .pluck(:urn)
+      @school_year_groups = school_year_groups
       @patient_count = patient_count
+      @nhs_numbers = Set.new
     end
 
     def self.call(...) = new(...).call
 
     def call
-      generate
       write_cohort_import_csv
-      write_class_import_csv
     end
 
     def patients
-      @patients = patient_count.times.map { build_patient }
+      patient_count.times.lazy.map { build_patient }
     end
 
     private
@@ -53,12 +77,6 @@ def cohort_import_csv_filepath
       )
     end
 
-    def class_import_csv_filepath(school:)
-      Rails.root.join(
-        "tmp/perf-test-class-import-#{school.name}-#{school.sessions.first.slug}.csv"
-      )
-    end
-
     def write_cohort_import_csv
       CSV.open(cohort_import_csv_filepath, "w") do |csv|
         csv << %w[
@@ -105,40 +123,7 @@ def write_cohort_import_csv
           ]
         end
       end
-    end
-
-    def write_class_import_csv
-      patients
-        .group_by(&:school)
-        .each do |school, school_patients|
-          next if school.nil?
-
-          CSV.open(class_import_csv_filepath(school:), "w") do |csv|
-            csv << %w[
-              CHILD_POSTCODE
-              CHILD_DATE_OF_BIRTH
-              CHILD_FIRST_NAME
-              CHILD_LAST_NAME
-              PARENT_1_EMAIL
-              PARENT_1_PHONE
-              PARENT_2_EMAIL
-              PARENT_2_PHONE
-            ]
-
-            school_patients.each do |patient|
-              csv << [
-                patient.address_postcode,
-                patient.date_of_birth,
-                patient.given_name,
-                patient.family_name,
-                patient.parents.first&.email,
-                patient.parents.first&.phone,
-                patient.parents.second&.email,
-                patient.parents.second&.phone
-              ]
-            end
-          end
-        end
+      cohort_import_csv_filepath.to_s
     end
 
     def programme_year_groups
@@ -147,28 +132,46 @@ def programme_year_groups
 
     def schools_with_year_groups
       @schools_with_year_groups ||=
-        organisation
-          .locations
-          .includes(:organisation, :sessions)
-          .select { (it.year_groups & programme_year_groups).any? }
+        begin
+          locations =
+            if school_year_groups.present?
+              urns.map do |urn|
+                Location.new(urn:, year_groups: school_year_groups[urn])
+              end
+            else
+              organisation
+                .locations
+                .where(urn: urns)
+                .includes(:organisation, :sessions)
+            end
+          locations.select { (it.year_groups & programme_year_groups).any? }
+        end
     end
 
-    def build_patient(year_group: nil)
+    def build_patient
       school = schools_with_year_groups.sample
       year_group ||= (school.year_groups & programme_year_groups).sample
+      nhs_number = nil
+      loop do
+        nhs_number = Faker::NationalHealthService.british_number.gsub(" ", "")
+        break unless nhs_number.in? @nhs_numbers
+      end
+      @nhs_numbers << nhs_number
 
       FactoryBot
         .build(
           :patient,
           school:,
-          date_of_birth: date_of_birth_for_year(year_group)
+          organisation:,
+          date_of_birth: date_of_birth_for_year(year_group),
+          nhs_number:
         )
         .tap do |patient|
           patient.parents =
             FactoryBot.build_list(:parent, 2, family_name: patient.family_name)
           patient.parent_relationships =
             patient.parents.map do
-              FactoryBot.build(:parent_relationship, parent: it)
+              FactoryBot.build(:parent_relationship, parent: it, patient:)
             end
         end
     end

lib/generate/consents.rb

@@ -12,6 +12,8 @@ def initialize(
       given: 0,
       given_needs_triage: 0
     )
+      validate_programme_and_session(programme, session) if programme
+
       @organisation = organisation
       @programme = programme || organisation.programmes.sample
       @session = session
@@ -24,18 +26,34 @@ def call
       create_consent_with_response(:refused, @refused)
       create_consent_with_response(:given, @given)
       create_consent_given_needs_triage(@given_needs_triage)
+      StatusUpdater.call(patient: patients)
     end
 
     def self.call(...) = new(...).call
 
     private
 
     def patients
-      (@session.presence || organisation)
-        .patients
-        .includes(:parents, :consents, consents: :parent)
-        .in_programmes([programme])
-        .select { it.consents.empty? && it.parents.any? }
+      @patients ||=
+        begin
+          sessions =
+            if @session
+              [@session]
+            else
+              organisation
+                .sessions
+                .eager_load(:location)
+                .merge(Location.school)
+                .has_programme(programme)
+            end
+
+          sessions.flat_map do |session|
+            session
+              .patients
+              .includes(:parents, :school, :consents, consents: :parent)
+              .select { it.consents.empty? && it.parents.any? }
+          end
+        end
     end
 
     def random_patients(count)
@@ -50,12 +68,13 @@ def random_patients(count)
     end
 
     def session_for(patient)
-      patient
-        .sessions
-        .eager_load(:location)
-        .merge(Location.school)
-        .has_programme(programme)
-        .sample
+      @session ||
+        patient
+          .sessions
+          .eager_load(:location)
+          .merge(Location.school)
+          .has_programme(programme)
+          .sample
     end
 
     def create_consent_with_response(response, count)
@@ -64,11 +83,13 @@ def create_consent_with_response(response, count)
 
       available_patient_sessions.each do |patient, session|
         consent = FactoryBot.create(:consent, response, patient:, programme:)
+        school = session.location.school? ? session.location : patient.school
         FactoryBot.create(
           :consent_form,
           organisation:,
           programmes: [programme],
           session:,
+          school:,
           consent:,
           response:
         )
@@ -98,5 +119,15 @@ def create_consent_given_needs_triage(count)
         )
       end
     end
+
+    def validate_programme_and_session(programme, session)
+      if session
+        if session.programmes.exclude?(programme)
+          raise "Session does not support programme #{programme.type}"
+        end
+      elsif programme.sessions.none? { it.location.school? }
+        raise "Programme #{programme.type} does not have a school session"
+      end
+    end
   end
 end