From 00a5f9cdfa2086df2b2db8ba914997f5e4920d6b Mon Sep 17 00:00:00 2001
From: Ryo Nakamura <r7kamura@gmail.com>
Date: Wed, 26 Oct 2022 06:22:11 +0900
Subject: [PATCH] Support `:traits` option on implicit association

---
 GETTING_STARTED.md                         |  2 ++
 lib/factory_bot/declaration/association.rb |  2 +-
 lib/factory_bot/definition_proxy.rb        |  8 ++++++-
 spec/acceptance/associations_spec.rb       | 28 ++++++++++++++++++++++
 spec/factory_bot/definition_proxy_spec.rb  |  9 +++++++
 5 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/GETTING_STARTED.md b/GETTING_STARTED.md
index 3be6d420f..94abfa77b 100644
--- a/GETTING_STARTED.md
+++ b/GETTING_STARTED.md
@@ -1296,6 +1296,8 @@ factory :post do
   association :author, :admin, factory: :user, name: 'John Doe'
   # or
   association :author, factory: [:user, :admin], name: 'John Doe'
+  # or
+  author factory: :user, traits: [:admin], name: 'John Doe'
 end
 
 # creates an admin user with name "John Doe"
diff --git a/lib/factory_bot/declaration/association.rb b/lib/factory_bot/declaration/association.rb
index 123e73b9e..91d4e4c9b 100644
--- a/lib/factory_bot/declaration/association.rb
+++ b/lib/factory_bot/declaration/association.rb
@@ -7,7 +7,7 @@ def initialize(name, *options)
         @options = options.dup
         @overrides = options.extract_options!
         @factory_name = @overrides.delete(:factory) || name
-        @traits = options
+        @traits = options + (@overrides.delete(:traits) || [])
       end
 
       def ==(other)
diff --git a/lib/factory_bot/definition_proxy.rb b/lib/factory_bot/definition_proxy.rb
index 5a6e4c025..b2d9a8af8 100644
--- a/lib/factory_bot/definition_proxy.rb
+++ b/lib/factory_bot/definition_proxy.rb
@@ -15,6 +15,11 @@ class DefinitionProxy
       method
     ].freeze
 
+    VALID_OPTION_KEYS = %i[
+      factory
+      traits
+    ].freeze
+
     (instance_methods + private_instance_methods).each do |method|
       undef_method(method) unless UNPROXIED_METHODS.include?(method.to_s)
     end
@@ -250,7 +255,8 @@ def __declare_attribute__(name, block)
     end
 
     def __valid_association_options?(options)
-      options.respond_to?(:has_key?) && options.has_key?(:factory)
+      options.respond_to?(:has_key?) &&
+        VALID_OPTION_KEYS.any? { |key| options.has_key?(key) }
     end
   end
 end
diff --git a/spec/acceptance/associations_spec.rb b/spec/acceptance/associations_spec.rb
index 6bf409faa..60412feb8 100644
--- a/spec/acceptance/associations_spec.rb
+++ b/spec/acceptance/associations_spec.rb
@@ -142,4 +142,32 @@
       )
     end
   end
+
+  context "when building implicit association with :traits option" do
+    it "builds the association according to the given strategy" do
+      define_model("Article", user_id: :integer) do
+        belongs_to :user
+      end
+
+      define_model("User") do
+        has_many :articles
+        attr_accessor :active
+      end
+
+      FactoryBot.define do
+        factory :article do
+          user traits: [:active]
+        end
+
+        factory :user do
+          trait :active do
+            active { true }
+          end
+        end
+      end
+
+      article = FactoryBot.create(:article)
+      expect(article.user.active).to be true
+    end
+  end
 end
diff --git a/spec/factory_bot/definition_proxy_spec.rb b/spec/factory_bot/definition_proxy_spec.rb
index 4523cd9fa..0596ceea6 100644
--- a/spec/factory_bot/definition_proxy_spec.rb
+++ b/spec/factory_bot/definition_proxy_spec.rb
@@ -53,6 +53,15 @@
       .with_options(factory: :user)
   end
 
+  it 'declares an association when called with a ":traits" key' do
+    definition = FactoryBot::Definition.new(:name)
+    proxy = FactoryBot::DefinitionProxy.new(definition)
+    proxy.author traits: [:admin]
+
+    expect(definition).to have_association_declaration(:author)
+      .with_options(traits: [:admin])
+  end
+
   it "declares a dynamic attribute when called with a block" do
     definition = FactoryBot::Definition.new(:name)
     proxy = FactoryBot::DefinitionProxy.new(definition)