Add functionality for enum traits

lib/factory_bot.rb

@@ -31,6 +31,7 @@
 require "factory_bot/sequence"
 require "factory_bot/attribute_list"
 require "factory_bot/trait"
+require "factory_bot/enum"
 require "factory_bot/aliases"
 require "factory_bot/definition"
 require "factory_bot/definition_proxy"
@@ -61,6 +62,9 @@ def self.reset_configuration
   mattr_accessor :use_parent_strategy, instance_accessor: false
   self.use_parent_strategy = true
 
+  mattr_accessor :automatically_define_enum_traits, instance_accessor: false
+  self.automatically_define_enum_traits = true
+
   # Look for errors in factories and (optionally) their traits.
   # Parameters:
   # factories - which factories to lint; omit for all factories

lib/factory_bot/definition.rb

@@ -1,13 +1,14 @@
 module FactoryBot
   # @api private
   class Definition
-    attr_reader :defined_traits, :declarations, :name
+    attr_reader :defined_traits, :declarations, :name, :registered_enums
 
     def initialize(name, base_traits = [])
       @name              = name
       @declarations      = DeclarationList.new(name)
       @callbacks         = []
       @defined_traits    = Set.new
+      @registered_enums  = []
       @to_create         = nil
       @base_traits       = base_traits
       @additional_traits = []
@@ -43,8 +44,10 @@ def callbacks
       aggregate_from_traits_and_self(:callbacks) { @callbacks }
     end
 
-    def compile
+    def compile(klass = nil)
       unless @compiled
+        expand_enum_traits(klass) unless klass.nil?
+
         declarations.attributes
 
         defined_traits.each do |defined_trait|
@@ -81,6 +84,10 @@ def define_trait(trait)
       @defined_traits.add(trait)
     end
 
+    def register_enum(enum)
+      @registered_enums << enum
+    end
+
     def define_constructor(&block)
       @constructor = block
     end
@@ -135,5 +142,25 @@ def aggregate_from_traits_and_self(method_name, &block)
         additional_traits.map(&method_name),
       ].flatten.compact
     end
+
+    def expand_enum_traits(klass)
+      if automatically_register_defined_enums?(klass)
+        automatically_register_defined_enums(klass)
+      end
+
+      registered_enums.each do |enum|
+        traits = enum.build_traits(klass)
+        traits.each { |trait| define_trait(trait) }
+      end
+    end
+
+    def automatically_register_defined_enums(klass)
+      klass.defined_enums.each_key { |name| register_enum(Enum.new(name)) }
+    end
+
+    def automatically_register_defined_enums?(klass)
+      FactoryBot.automatically_define_enum_traits &&
+        klass.respond_to?(:defined_enums)
+    end
   end
 end

lib/factory_bot/definition_proxy.rb

@@ -176,6 +176,64 @@ def trait(name, &block)
       @definition.define_trait(Trait.new(name, &block))
     end
 
+    # Creates traits for enumerable values.
+    #
+    # Example:
+    #   factory :task do
+    #     traits_for_enum :status, [:started, :finished]
+    #   end
+    #
+    # Equivalent to:
+    #   factory :task do
+    #     trait :started do
+    #       status { :started }
+    #     end
+    #
+    #     trait :finished do
+    #       status { :finished }
+    #     end
+    #   end
+    #
+    # Example:
+    #   factory :task do
+    #     traits_for_enum :status, {started: 1, finished: 2}
+    #   end
+    #
+    # Example:
+    #   class Task
+    #     def statuses
+    #       {started: 1, finished: 2}
+    #     end
+    #   end
+    #
+    #   factory :task do
+    #     traits_for_enum :status
+    #   end
+    #
+    # Both equivalent to:
+    #   factory :task do
+    #     trait :started do
+    #       status { 1 }
+    #     end
+    #
+    #     trait :finished do
+    #       status { 2 }
+    #     end
+    #   end
+    #
+    #
+    # Arguments:
+    #   attribute_name: +Symbol+ or +String+
+    #     the name of the attribute these traits will set the value of
+    #   values: +Array+, +Hash+, or other +Enumerable+
+    #     An array of trait names, or a mapping of trait names to values for
+    #     those traits. When this argument is not provided, factory_bot will
+    #     attempt to get the values by calling the pluralized `attribute_name`
+    #     class method.
+    def traits_for_enum(attribute_name, values = nil)
+      @definition.register_enum(Enum.new(attribute_name, values))
+    end
+
     def initialize_with(&block)
       @definition.define_constructor(&block)
     end

lib/factory_bot/enum.rb

@@ -0,0 +1,27 @@
+module FactoryBot
+  # @api private
+  class Enum
+    def initialize(attribute_name, values = nil)
+      @attribute_name = attribute_name
+      @values = values
+    end
+
+    def build_traits(klass)
+      enum_values(klass).map do |trait_name, value|
+        build_trait(trait_name, @attribute_name, value || trait_name)
+      end
+    end
+
+    private
+
+    def enum_values(klass)
+      @values || klass.send(@attribute_name.to_s.pluralize)
+    end
+
+    def build_trait(trait_name, attribute_name, value)
+      Trait.new(trait_name) do
+        add_attribute(attribute_name) { value }
+      end
+    end
+  end
+end

lib/factory_bot/factory.rb

@@ -84,7 +84,7 @@ def compile
       unless @compiled
         parent.compile
         parent.defined_traits.each { |trait| define_trait(trait) }
-        @definition.compile
+        @definition.compile(build_class)
         build_hierarchy
         @compiled = true
       end

spec/acceptance/enum_traits_spec.rb

@@ -0,0 +1,161 @@
+describe "enum traits" do
+  context "when automatically_define_enum_traits is true" do
+    it "builds traits automatically for model enum field" do
+      define_model("Task", status: :integer) do
+        enum status: { queued: 0, started: 1, finished: 2 }
+      end
+
+      FactoryBot.define do
+        factory :task
+      end
+
+      Task.statuses.each_key do |trait_name|
+        task = FactoryBot.build(:task, trait_name)
+
+        expect(task.status).to eq(trait_name)
+      end
+
+      Task.reset_column_information
+    end
+
+    it "prefers user defined traits over automatically built traits" do
+      define_model("Task", status: :integer) do
+        enum status: { queued: 0, started: 1, finished: 2 }
+      end
+
+      FactoryBot.define do
+        factory :task do
+          trait :queued do
+            status { :finished }
+          end
+
+          trait :started do
+            status { :finished }
+          end
+
+          trait :finished do
+            status { :finished }
+          end
+        end
+      end
+
+      Task.statuses.each_key do |trait_name|
+        task = FactoryBot.build(:task, trait_name)
+
+        expect(task.status).to eq("finished")
+      end
+
+      Task.reset_column_information
+    end
+
+    it "builds traits for each enumerated value using a provided list of values as a Hash" do
+      statuses = { queued: 0, started: 1, finished: 2 }
+
+      define_class "Task" do
+        attr_accessor :status
+      end
+
+      FactoryBot.define do
+        factory :task do
+          traits_for_enum :status, statuses
+        end
+      end
+
+      statuses.each do |trait_name, trait_value|
+        task = FactoryBot.build(:task, trait_name)
+
+        expect(task.status).to eq(trait_value)
+      end
+    end
+
+    it "builds traits for each enumerated value using a provided list of values as an Array" do
+      statuses = %w[queued started finished]
+
+      define_class "Task" do
+        attr_accessor :status
+      end
+
+      FactoryBot.define do
+        factory :task do
+          traits_for_enum :status, statuses
+        end
+      end
+
+      statuses.each do |trait_name|
+        task = FactoryBot.build(:task, trait_name)
+
+        expect(task.status).to eq(trait_name)
+      end
+    end
+
+    it "builds traits for each enumerated value using a custom enumerable" do
+      statuses = define_class("Statuses") do
+        include Enumerable
+
+        def each(&block)
+          ["queued", "started", "finished"].each(&block)
+        end
+      end.new
+
+      define_class "Task" do
+        attr_accessor :status
+      end
+
+      FactoryBot.define do
+        factory :task do
+          traits_for_enum :status, statuses
+        end
+      end
+
+      statuses.each do |trait_name|
+        task = FactoryBot.build(:task, trait_name)
+
+        expect(task.status).to eq(trait_name)
+      end
+    end
+  end
+
+  context "when automatically_define_enum_traits is false" do
+    it "raises an error for undefined traits" do
+      with_temporary_assignment(FactoryBot, :automatically_define_enum_traits, false) do
+        define_model("Task", status: :integer) do
+          enum status: { queued: 0, started: 1, finished: 2 }
+        end
+
+        FactoryBot.define do
+          factory :task
+        end
+
+        Task.statuses.each_key do |trait_name|
+          expect { FactoryBot.build(:task, trait_name) }.to raise_error(
+            KeyError, "Trait not registered: \"#{trait_name}\""
+          )
+        end
+
+        Task.reset_column_information
+      end
+    end
+
+    it "builds traits for each enumerated value when traits_for_enum are specified" do
+      with_temporary_assignment(FactoryBot, :automatically_define_enum_traits, false) do
+        define_model("Task", status: :integer) do
+          enum status: { queued: 0, started: 1, finished: 2 }
+        end
+
+        FactoryBot.define do
+          factory :task do
+            traits_for_enum(:status)
+          end
+        end
+
+        Task.statuses.each_key do |trait_name|
+          task = FactoryBot.build(:task, trait_name)
+
+          expect(task.status).to eq(trait_name)
+        end
+
+        Task.reset_column_information
+      end
+    end
+  end
+end

spec/factory_bot/definition_spec.rb

@@ -72,4 +72,14 @@
 
     expect(definition.to_create).to eq block
   end
+
+  it "maintains a list of enum fields" do
+    definition = described_class.new(:name)
+
+    enum_field = double("enum_field")
+
+    definition.register_enum(enum_field)
+
+    expect(definition.registered_enums).to include(enum_field)
+  end
 end

spec/factory_bot/factory_spec.rb

@@ -17,6 +17,7 @@
   end
 
   it "returns associations" do
+    define_class("Post")
     factory = FactoryBot::Factory.new(:post)
     FactoryBot::Internal.register_factory(FactoryBot::Factory.new(:admin))
     factory.declare_attribute(FactoryBot::Declaration::Association.new(:author, {}))
@@ -32,6 +33,7 @@
     association_on_parent = FactoryBot::Declaration::Association.new(:association_on_parent, {})
     association_on_child  = FactoryBot::Declaration::Association.new(:association_on_child, {})
 
+    define_class("Post")
     factory = FactoryBot::Factory.new(:post)
     factory.declare_attribute(association_on_parent)
     FactoryBot::Internal.register_factory(factory)
@@ -134,6 +136,7 @@
 
   it "creates a new factory while overriding the parent class" do
     name    = :user
+    define_class("User")
     factory = FactoryBot::Factory.new(name)
     FactoryBot::Internal.register_factory(factory)