From c9ed3d88e5aa8d02ffd3e8b63ea80542caee1a01 Mon Sep 17 00:00:00 2001
From: Joshua Clayton <jclayton@thoughtbot.com>
Date: Thu, 2 Aug 2012 10:19:17 -0400
Subject: [PATCH] Improve ability to define associations with traits and add
 documentation

---
 GETTING_STARTED.md                          | 38 +++++++++++++++++++++
 lib/factory_girl/declaration/association.rb | 10 +++---
 lib/factory_girl/definition_proxy.rb        |  4 +--
 spec/acceptance/traits_spec.rb              | 29 ++++++++++++++++
 4 files changed, 75 insertions(+), 6 deletions(-)

diff --git a/GETTING_STARTED.md b/GETTING_STARTED.md
index 3280c59..51122ab 100644
--- a/GETTING_STARTED.md
+++ b/GETTING_STARTED.md
@@ -632,7 +632,45 @@ end
 FactoryGirl.create_list(:user, 3, :admin, :male, name: "Jon Snow")
 ```
 
+Traits can be used with associations easily too:
 
+```ruby
+factory :user do
+  name "Friendly User"
+
+  trait :admin do
+    admin true
+  end
+end
+
+factory :post do
+  association :user, :admin, name: 'John Doe'
+end
+
+# creates an admin user with named "John Doe"
+FactoryGirl.create(:post).user
+```
+
+When you're using association names that're different than the factory:
+
+```ruby
+factory :user do
+  name "Friendly User"
+
+  trait :admin do
+    admin true
+  end
+end
+
+factory :post do
+  association :author, :admin, factory: :user, name: 'John Doe'
+  # or
+  association :author, factory: [:user, :admin], name: 'John Doe'
+end
+
+# creates an admin user with named "John Doe"
+FactoryGirl.create(:post).user
+```
 Callbacks
 ---------
 
diff --git a/lib/factory_girl/declaration/association.rb b/lib/factory_girl/declaration/association.rb
index 5588bcf..2821e34 100644
--- a/lib/factory_girl/declaration/association.rb
+++ b/lib/factory_girl/declaration/association.rb
@@ -2,9 +2,11 @@ module FactoryGirl
   class Declaration
     # @api private
     class Association < Declaration
-      def initialize(name, options)
+      def initialize(name, *options)
         super(name, false)
-        @options = options
+        @options = options.dup
+        @overrides = options.extract_options!
+        @traits = options
       end
 
       def ==(other)
@@ -18,8 +20,8 @@ module FactoryGirl
       private
 
       def build
-        factory_name = @options[:factory] || name
-        [Attribute::Association.new(name, factory_name, @options.except(:factory))]
+        factory_name = @overrides[:factory] || name
+        [Attribute::Association.new(name, factory_name, [@traits, @overrides.except(:factory)].flatten)]
       end
     end
   end
diff --git a/lib/factory_girl/definition_proxy.rb b/lib/factory_girl/definition_proxy.rb
index 4fb89d1..5e5d835 100644
--- a/lib/factory_girl/definition_proxy.rb
+++ b/lib/factory_girl/definition_proxy.rb
@@ -138,8 +138,8 @@ module FactoryGirl
     #    If no name is given, the name of the attribute is assumed to be the
     #    name of the factory. For example, a "user" association will by
     #    default use the "user" factory.
-    def association(name, options = {})
-      @definition.declare_attribute(Declaration::Association.new(name, options))
+    def association(name, *options)
+      @definition.declare_attribute(Declaration::Association.new(name, *options))
     end
 
     def to_create(&block)
diff --git a/spec/acceptance/traits_spec.rb b/spec/acceptance/traits_spec.rb
index aff6b7e..4fb2e0c 100644
--- a/spec/acceptance/traits_spec.rb
+++ b/spec/acceptance/traits_spec.rb
@@ -673,6 +673,15 @@ end
 describe "traits used in associations" do
   before do
     define_model("User", admin: :boolean, name: :string)
+
+    define_model("Comment", user_id: :integer) do
+      belongs_to :user
+    end
+
+    define_model("Order", creator_id: :integer) do
+      belongs_to :creator, class_name: 'User'
+    end
+
     define_model("Post", author_id: :integer) do
       belongs_to :author, class_name: 'User'
     end
@@ -689,6 +698,14 @@ describe "traits used in associations" do
       factory :post do
         association :author, factory: [:user, :admin], name: 'John Doe'
       end
+
+      factory :comment do
+        association :user, :admin, name: 'Joe Slick'
+      end
+
+      factory :order do
+        association :creator, :admin, factory: :user, name: 'Joe Creator'
+      end
     end
   end
 
@@ -697,4 +714,16 @@ describe "traits used in associations" do
     author.should be_admin
     author.name.should == 'John Doe'
   end
+
+  it "allows inline traits with the default association" do
+    user = FactoryGirl.create(:comment).user
+    user.should be_admin
+    user.name.should == 'Joe Slick'
+  end
+
+  it "allows inline traits with a specific factory for an association" do
+    creator = FactoryGirl.create(:order).creator
+    creator.should be_admin
+    creator.name.should == 'Joe Creator'
+  end
 end