Make spec/lib/gitlab/import_export/all_models.yml EE-ready

This port https://gitlab.com/gitlab-org/gitlab-ee/merge_requests/7254 to
CE.

Signed-off-by: Rémy Coutable <remy@rymai.me>

1 files changed, 26 insertions, 16 deletions

diff --git a/spec/lib/gitlab/import_export/model_configuration_spec.rb b/spec/lib/gitlab/import_export/model_configuration_spec.rb
index 5cb8f2589c8..2e28f978c3a 100644
--- a/spec/lib/gitlab/import_export/model_configuration_spec.rb
+++ b/spec/lib/gitlab/import_export/model_configuration_spec.rb
@@ -16,14 +16,30 @@ describe 'Import/Export model configuration' do
     # - User, Author... Models we do not care about for checking models
     names.flatten.uniq - %w(milestones labels user author) + ['project']
   end
+  let(:ce_models_yml) { 'spec/lib/gitlab/import_export/all_models.yml' }
+  let(:ce_models_hash) { YAML.load_file(ce_models_yml) }
+
+  let(:ee_models_yml) { 'ee/spec/lib/gitlab/import_export/all_models.yml' }
+  let(:ee_models_hash) { File.exist?(ee_models_yml) ? YAML.load_file(ee_models_yml) : {} }
 
-  let(:all_models_yml) { 'spec/lib/gitlab/import_export/all_models.yml' }
-  let(:all_models) { YAML.load_file(all_models_yml) }
   let(:current_models) { setup_models }
+  let(:all_models_hash) do
+    all_models_hash = ce_models_hash.dup
+
+    all_models_hash.each do |model, associations|
+      associations.concat(ee_models_hash[model] || [])
+    end
+
+    ee_models_hash.each do |model, associations|
+      all_models_hash[model] ||= associations
+    end
+
+    all_models_hash
+  end
 
   it 'has no new models' do
     model_names.each do |model_name|
-      new_models = Array(current_models[model_name]) - Array(all_models[model_name])
+      new_models = Array(current_models[model_name]) - Array(all_models_hash[model_name])
       expect(new_models).to be_empty, failure_message(model_name.classify, new_models)
     end
   end
@@ -31,27 +47,21 @@ describe 'Import/Export model configuration' do
   # List of current models between models, in the format of
   # {model: [model_2, model3], ...}
   def setup_models
-    all_models_hash = {}
-
-    model_names.each do |model_name|
-      model_class = relation_class_for_name(model_name)
-
-      all_models_hash[model_name] = associations_for(model_class) - ['project']
+    model_names.each_with_object({}) do |model_name, hash|
+      hash[model_name] = associations_for(relation_class_for_name(model_name)) - ['project']
     end
-
-    all_models_hash
   end
 
   def failure_message(parent_model_name, new_models)
-    <<-MSG
+    <<~MSG
       New model(s) <#{new_models.join(',')}> have been added, related to #{parent_model_name}, which is exported by
       the Import/Export feature.
 
-      If you think this model should be included in the export, please add it to IMPORT_EXPORT_CONFIG.
-      Definitely add it to MODELS_JSON to signal that you've handled this error and to prevent it from showing up in the future.
+      If you think this model should be included in the export, please add it to `#{Gitlab::ImportExport.config_file}`.
 
-      MODELS_JSON: #{File.expand_path(all_models_yml)}
-      IMPORT_EXPORT_CONFIG: #{Gitlab::ImportExport.config_file}
+      Definitely add it to `#{File.expand_path(ce_models_yml)}`
+      #{"or `#{File.expand_path(ee_models_yml)}` if the model/associations are EE-specific\n" if ee_models_hash.any?}
+      to signal that you've handled this error and to prevent it from showing up in the future.
     MSG
   end
 end