Add polymorphic association support

Author: tiagolnobre

README.md

@@ -238,6 +238,46 @@ FixtureBot.define do
 end
 ```
 
+### Polymorphic associations
+
+Reference records using Rails polymorphic associations:
+
+```ruby
+FixtureBot.define do
+  post :hello_world do
+    title "Hello World"
+    author :brad
+  end
+
+  user :alice do
+    name "Alice"
+  end
+
+  vote :like do
+    user :alice
+    votable post(:hello_world)  # sets votable_id and votable_type: Post
+  end
+end
+```
+
+The `votable post(:hello_world)` syntax sets both `votable_id` (to the post's stable ID) and `votable_type` (to `"Post"`).
+
+FixtureBot auto-detects polymorphic associations from your schema columns. Just include both `_id` and `_type` columns:
+
+```ruby
+FixtureBot::Schema.define do
+  table :posts, singular: :post, columns: [:title, :author_id] do
+    belongs_to :author, table: :users
+  end
+
+  table :votes, singular: :vote, columns: [:user_id, :votable_id, :votable_type] do
+    belongs_to :user, table: :users
+  end
+end
+```
+
+In Rails, FixtureBot auto-detects polymorphic columns (`_id` + `_type` pairs) from your database schema automatically.
+
 ### Implicit vs explicit style
 
 By default, the block is evaluated implicitly. Table methods like `user` and `post` are available directly:

lib/fixturebot/definition.rb

@@ -21,8 +21,10 @@ def define_table_method(table)
       define_singleton_method(table.singular_name) do |record_name = nil, &block|
         if record_name.nil? && block.nil?
           Default::Definition.new(table, @defaults[table.name])
-        elsif record_name
+        elsif record_name && block
           add_row(table, record_name, block)
+        elsif record_name
+          add_row(table, record_name, nil)
         else
           raise ArgumentError, "#{table.singular_name} requires a record name or no arguments"
         end
@@ -37,7 +39,8 @@ def add_row(table, record_name, block)
         name: record_name,
         literal_values: row_def.literal_values,
         association_refs: row_def.association_refs,
-        tag_refs: row_def.tag_refs
+        tag_refs: row_def.tag_refs,
+        polymorphic_refs: row_def.polymorphic_refs
       )
     end
   end

lib/fixturebot/rails/schema_loader.rb

@@ -42,22 +42,46 @@ def build_table(name)
           .reject { |c| framework_column?(c.name) }
           .map { |c| c.name.to_sym }
 
+        column_names = columns.map(&:to_s)
+        polymorphic_pairs = detect_polymorphic_columns(column_names)
+
         associations = @connection.foreign_keys(name).map do |fk|
+          next if polymorphic_pairs.key?(fk.column.to_s)
+
           Schema::BelongsTo.new(
             name: association_name(fk.column),
             table: fk.to_table.to_sym,
             foreign_key: fk.column.to_sym
           )
+        end.compact
+
+        polymorphic_associations = polymorphic_pairs.map do |id_col, type_col|
+          Schema::PolymorphicBelongsTo.new(
+            name: association_name(id_col),
+            foreign_key: id_col.to_sym,
+            type_column: type_col.to_sym
+          )
         end
 
         Schema::Table.new(
           name: name.to_sym,
           singular_name: singularize(name),
           columns: columns,
-          belongs_to_associations: associations
+          belongs_to_associations: associations,
+          polymorphic_belongs_to_associations: polymorphic_associations
         )
       end
 
+      def detect_polymorphic_columns(column_names)
+        pairs = {}
+        id_columns = column_names.select { |c| c.end_with?("_id") }
+        id_columns.each do |id_col|
+          type_col = id_col.sub(/_id$/, "_type")
+          pairs[id_col] = type_col if column_names.include?(type_col)
+        end
+        pairs
+      end
+
       def build_join_table(name)
         fk_columns = foreign_key_columns(name)
 

lib/fixturebot/row.rb

@@ -2,19 +2,33 @@
 
 module FixtureBot
   module Row
-    Declaration = Data.define(:table, :name, :literal_values, :association_refs, :tag_refs)
+    TableReference = Data.define(:table_name, :record_name)
+    Declaration = Data.define(:table, :name, :literal_values, :association_refs, :tag_refs, :polymorphic_refs)
 
     class Definition
-      attr_reader :literal_values, :association_refs, :tag_refs
+      attr_reader :literal_values, :association_refs, :tag_refs, :polymorphic_refs
 
       def initialize(table, schema)
         @literal_values = {}
         @association_refs = {}
         @tag_refs = {}
+        @polymorphic_refs = {}
 
         define_column_methods(table)
         define_association_methods(table)
+        define_polymorphic_methods(table)
         define_join_table_methods(table, schema)
+        define_table_reference_methods(schema, table)
+      end
+
+      def define_table_reference_methods(schema, table)
+        schema.tables.each_value do |tbl|
+          next if table.belongs_to_associations.any? { |a| a.name == tbl.singular_name }
+          next if table.polymorphic_belongs_to_associations.any? { |a| a.name == tbl.singular_name }
+          define_singleton_method(tbl.singular_name) do |record_name|
+            TableReference.new(table_name: tbl.name, record_name: record_name)
+          end
+        end
       end
 
       private
@@ -35,6 +49,14 @@ def define_association_methods(table)
         end
       end
 
+      def define_polymorphic_methods(table)
+        table.polymorphic_belongs_to_associations.each do |assoc|
+          define_singleton_method(assoc.name) do |table_ref|
+            @polymorphic_refs[assoc.name] = table_ref
+          end
+        end
+      end
+
       def define_join_table_methods(table, schema)
         schema.join_tables.each_value do |jt|
           if jt.left_table == table.name
@@ -69,6 +91,8 @@ def record
             result[col] = @row.literal_values[col]
           elsif foreign_key_values.key?(col)
             result[col] = foreign_key_values[col]
+          elsif polymorphic_values.key?(col)
+            result[col] = polymorphic_values[col]
           elsif defaulted_values.key?(col)
             result[col] = defaulted_values[col]
           end
@@ -112,10 +136,27 @@ def foreign_key_values
         end
       end
 
+      def polymorphic_values
+        @polymorphic_values ||= @row.polymorphic_refs.each_with_object({}) do |(assoc_name, table_ref), hash|
+          assoc = @table.polymorphic_belongs_to_associations.find { |a| a.name == assoc_name }
+          hash[assoc.foreign_key] = Key.generate(table_ref.table_name, table_ref.record_name)
+          hash[assoc.type_column] = classify(table_ref.table_name)
+        end
+      end
+
+      def classify(table_name)
+        if defined?(ActiveSupport::Inflector)
+          ActiveSupport::Inflector.classify(table_name.to_s)
+        else
+          table_name.to_s.sub(/s$/, "").capitalize
+        end
+      end
+
       def defaulted_values
         @defaulted_values ||= @defaults.each_with_object({}) do |(col, block), result|
           next if @row.literal_values.key?(col)
           next if foreign_key_values.key?(col)
+          next if polymorphic_values.key?(col)
 
           fixture = Default::Fixture.new(key: @row.name)
           context = Default::Context.new(literal_values: @row.literal_values)

lib/fixturebot/schema.rb

@@ -2,8 +2,9 @@
 
 module FixtureBot
   class Schema
-    Table = Data.define(:name, :singular_name, :columns, :belongs_to_associations)
+    Table = Data.define(:name, :singular_name, :columns, :belongs_to_associations, :polymorphic_belongs_to_associations)
     BelongsTo = Data.define(:name, :table, :foreign_key)
+    PolymorphicBelongsTo = Data.define(:name, :foreign_key, :type_column)
     JoinTable = Data.define(:name, :left_table, :right_table, :left_foreign_key, :right_foreign_key)
 
     attr_reader :tables, :join_tables
@@ -35,15 +36,33 @@ def initialize(schema)
 
       def table(name, singular:, columns: [], &block)
         associations = []
+        polymorphic_associations = []
         if block
-          table_builder = TableBuilder.new(associations)
+          table_builder = TableBuilder.new(associations, polymorphic_associations, columns)
           table_builder.instance_eval(&block)
         end
+
+        columns_set = columns.to_set
+        columns_set.grep(/_id$/).each do |id_col|
+          type_col = id_col.to_s.sub(/_id$/, "_type").to_sym
+          if columns_set.include?(type_col)
+            assoc_name = id_col.to_s.sub(/_id$/, "").to_sym
+            next if associations.any? { |a| a.name == assoc_name }
+            next if polymorphic_associations.any? { |a| a.name == assoc_name }
+            polymorphic_associations << PolymorphicBelongsTo.new(
+              name: assoc_name,
+              foreign_key: id_col,
+              type_column: type_col
+            )
+          end
+        end
+
         @schema.add_table(Table.new(
           name: name,
           singular_name: singular,
           columns: columns,
-          belongs_to_associations: associations
+          belongs_to_associations: associations,
+          polymorphic_belongs_to_associations: polymorphic_associations
         ))
       end
 
@@ -61,13 +80,27 @@ def join_table(name, left_table, right_table)
     end
 
     class TableBuilder
-      def initialize(associations)
+      def initialize(associations, polymorphic_associations, columns)
         @associations = associations
+        @polymorphic_associations = polymorphic_associations
+        @columns = columns
       end
 
-      def belongs_to(name, table:)
+      def belongs_to(name, table: nil)
         foreign_key = :"#{name}_id"
-        @associations << BelongsTo.new(name: name, table: table, foreign_key: foreign_key)
+        type_column = :"#{name}_type"
+
+        if table
+          @associations << BelongsTo.new(name: name, table: table, foreign_key: foreign_key)
+        elsif @columns.include?(foreign_key) && @columns.include?(type_column)
+          @polymorphic_associations << PolymorphicBelongsTo.new(
+            name: name,
+            foreign_key: foreign_key,
+            type_column: type_column
+          )
+        else
+          raise ArgumentError, "belongs_to :#{name} requires table: option"
+        end
       end
     end
   end

playground/blog/fixtures.rb

@@ -51,4 +51,19 @@
     post :tdd_guide
     author :brad
   end
+
+  vote :hello_like do
+    user :alice
+    votable post(:hello_world)
+  end
+
+  vote :hello_upvote do
+    user :brad
+    votable post(:hello_world)
+  end
+
+  vote :tdd_upvote do
+    user :charlie
+    votable post(:tdd_guide)
+  end
 end

playground/blog/schema.rb

@@ -10,6 +10,10 @@
     belongs_to :author, table: :users
   end
 
+  table :votes, singular: :vote, columns: [:user_id, :votable_id, :votable_type] do
+    belongs_to :user, table: :users
+  end
+
   table :tags, singular: :tag, columns: [:name]
 
   join_table :posts_tags, :posts, :tags

spec/fixturebot/schema_loader_spec.rb

@@ -31,6 +31,22 @@
         t.integer "post_id", null: false
         t.integer "tag_id", null: false
       end
+
+      create_table "votes", force: :cascade do |t|
+        t.integer "user_id"
+        t.integer "votable_id"
+        t.string "votable_type"
+        t.timestamps
+      end
+
+      create_table "comments", force: :cascade do |t|
+        t.text "body"
+        t.integer "author_id"
+        t.string "author_type"
+        t.integer "parent_id"
+        t.string "parent_type"
+        t.timestamps
+      end
     end
   end
 
@@ -41,7 +57,7 @@
   subject(:schema) { described_class.load }
 
   it "loads regular tables with columns" do
-    expect(schema.tables.keys).to contain_exactly(:users, :posts, :tags)
+    expect(schema.tables.keys).to contain_exactly(:users, :posts, :tags, :votes, :comments)
   end
 
   it "skips id, created_at, updated_at columns" do
@@ -70,4 +86,27 @@
   it "does not include join tables in regular tables" do
     expect(schema.tables).not_to have_key(:posts_tags)
   end
+
+  it "detects polymorphic associations from _id and _type columns" do
+    polymorphic_assocs = schema.tables[:votes].polymorphic_belongs_to_associations
+    expect(polymorphic_assocs.size).to eq(1)
+
+    votable = polymorphic_assocs.first
+    expect(votable.name).to eq(:votable)
+    expect(votable.foreign_key).to eq(:votable_id)
+    expect(votable.type_column).to eq(:votable_type)
+  end
+
+  it "excludes polymorphic foreign keys from belongs_to associations" do
+    belongs_to = schema.tables[:votes].belongs_to_associations
+    expect(belongs_to).to be_empty
+  end
+
+  it "detects multiple polymorphic associations" do
+    polymorphic_assocs = schema.tables[:comments].polymorphic_belongs_to_associations
+    expect(polymorphic_assocs.size).to eq(2)
+
+    names = polymorphic_assocs.map(&:name)
+    expect(names).to include(:author, :parent)
+  end
 end