Keshav/polymorphic generator

lib/generators/seedie/install_generator.rb

@@ -1,9 +1,11 @@
 require "rails/generators/base"
-require "active_record"
+require "seedie"
 
 module Seedie
   module Generators
     class InstallGenerator < Rails::Generators::Base
+      include PolymorphicAssociationHelper
+
       EXCLUDED_MODELS = %w[
         ActiveRecord::SchemaMigration
         ActiveRecord::InternalMetadata
@@ -24,6 +26,7 @@ def generate_seedie_file(output = STDOUT)
         @models = get_models
         @models_config = build_models_config
         template "seedie.yml", "config/seedie.yml"
+
         output_seedie_warning(output)
       end
 
@@ -100,14 +103,24 @@ def associations_configuration(model)
 
       def belongs_to_associations_configuration(model)
         belongs_to_associations = model.reflect_on_all_associations(:belongs_to).reject do |association|
-          association.options[:polymorphic] == true || # Excluded Polymorphic Associations
           association.options[:optional] == true # Excluded Optional Associations
         end
 
         belongs_to_associations.reduce({}) do |config, association|
-          config[association.name.to_s] = "random"
+          if association.polymorphic?
+            config[association.name.to_s] = set_polymorphic_association_config(model, association)
+          else
+            config[association.name.to_s] = "random"
+          end
           config
         end
+      end      
+
+      def set_polymorphic_association_config(model, association)
+        {
+          "polymorphic" => find_polymorphic_types(model, association.name),
+          "strategy" => "random"
+        }
       end
 
       def has_presence_validator?(model, column_name)

lib/generators/seedie/templates/seedie.yml

@@ -19,8 +19,15 @@ models:
 <% config['associations'].each do |type, associations| -%>
       <%= type %>:
 <% associations.each do |association, value| -%>
+<% if value.is_a?(Hash) && value.key?("polymorphic") -%>
+<% next if value["polymorphic"].blank? -%>
+        <%= association %>:
+          polymorphic: <%= value["polymorphic"] %>
+          strategy: <%= value["strategy"] %>
+<% else -%>
         <%= association %>: <%= value %>
 <% end -%>
 <% end -%>
 <% end -%>
 <% end -%>
+<% end -%>

lib/seedie.rb

@@ -11,6 +11,8 @@
 require_relative "seedie/model_fields"
 require_relative "seedie/model_seeder"
 
+require_relative "seedie/polymorphic_association_helper"
+
 require_relative "seedie/model/creator"
 require_relative "seedie/model/model_sorter"
 require_relative "seedie/model/id_generator"

lib/seedie/model/model_sorter.rb

@@ -1,6 +1,8 @@
 module Seedie
   module Model
     class ModelSorter
+      include PolymorphicAssociationHelper
+
       def initialize(models)
         @models = models
         @model_dependencies = models.map {|m| [m, get_model_dependencies(m)]}.to_h
@@ -51,7 +53,6 @@ def resolve_dependencies(model)
 
       def get_model_dependencies(model)
         associations = model.reflect_on_all_associations(:belongs_to).reject do |association|
-          association.options[:polymorphic] == true || # Excluded Polymorphic Associations
           association.options[:optional] == true # Excluded Optional Associations
         end
 
@@ -60,10 +61,17 @@ def get_model_dependencies(model)
         associations.map do |association|
           if association.options[:class_name]
             constantize_class_name(association.options[:class_name], model.name)
+          elsif association.polymorphic?
+            types = find_polymorphic_types(model, association.name)
+
+            if types.blank?
+              puts "Polymorphic type not found for #{model.name}. Ignoring..." 
+              next
+            end
           else
             association.klass
           end
-        end
+        end.compact
       end
 
       private

lib/seedie/polymorphic_association_helper.rb

@@ -0,0 +1,20 @@
+module PolymorphicAssociationHelper
+  # Returns the type of the polymorphic association
+  # We need only one polymorphic association while generating config
+  # this makes it easier to sort according to dependencies
+  def find_polymorphic_types(model, association_name)
+    type = @models.find { |potential_model| has_association?(potential_model, association_name) }
+    type&.name&.underscore
+  end
+
+  def has_association?(model, association_name)
+    associations = select_associations(model)
+    associations.any? { |association| association.options[:as] == association_name }
+  end
+
+  def select_associations(model)
+    model.reflect_on_all_associations.select do |reflection|
+      %i[has_many has_one].include?(reflection.macro)
+    end
+  end
+end

spec/generators/install_generator_spec.rb

@@ -81,22 +81,22 @@
   end
 
   describe "#belongs_to_associations_configuration" do
-    it "excludes polymorphic associations" do
-      expect(content["models"]["review"]["associations"]["belongs_to"]).not_to include("reviewable")
-    end
-
     it "excludes optional associations" do
       expect(content["models"]["comment"]["associations"]["belongs_to"]).not_to include("user")
     end
 
     it "generates required belongs_to associations" do
       expect(content["models"]["user"]["associations"]["belongs_to"]).to be nil
       expect(content["models"]["post"]["associations"]["belongs_to"]).to be nil
-      expect(content["models"]["review"]["associations"]["belongs_to"]).to include("user")
       expect(content["models"]["post_metadatum"]["associations"]["belongs_to"]).to include("post")
       expect(content["models"]["game_room"]["associations"]["belongs_to"]).to include("creator", "updater")
       expect(content["models"]["comment"]["associations"]["belongs_to"]).to include("post")
     end
+
+    it "generates required polymorphic belongs_to associations" do
+      expect(content["models"]["review"]["associations"]["belongs_to"]).to include("reviewable")
+      expect(content["models"]["review"]["associations"]["belongs_to"]["reviewable"]).to include("polymorphic")
+    end
   end
 
   describe "Validations" do