From d05a9a3c4c566d835dff6d946c91ddd86d96b28a Mon Sep 17 00:00:00 2001
From: Daniel Colson <danieljamescolson@gmail.com>
Date: Fri, 10 Jul 2020 11:53:20 -0400
Subject: [PATCH] Add definition names to default trait key errors (#1421)

* Add definition names to default trait key errors

Closes #1222

Before this commit, referencing a trait that didn't exist would raise a
generic `KeyError: Trait not registered: "trait_name"`. This can
sometime make it difficult to know where exactly the error is coming
from. The fact that implicitly declared associations and sequences will
fall back to implicit traits if they can't be found compounds this
problem. If various associations, sequences, or traits share the same
name, the hunt for the error can take some time.

With this commit we include the factory or trait name (i.e. the
definition name) in the error message to help identify where the
problematic trait originated.

Because trait lookup relies on a more generic class that raises a
`KeyError` for missing keys, we rescue that error, then construct a new
error with our custom message using the original error's message and
backtrace.
---
 lib/factory_bot/definition.rb  | 14 ++++++
 spec/acceptance/traits_spec.rb | 91 +++++++++++++++++++++++++++-------
 2 files changed, 86 insertions(+), 19 deletions(-)

diff --git a/lib/factory_bot/definition.rb b/lib/factory_bot/definition.rb
index 26ed835..4aded74 100644
--- a/lib/factory_bot/definition.rb
+++ b/lib/factory_bot/definition.rb
@@ -111,6 +111,20 @@ module FactoryBot
 
     def base_traits
       @base_traits.map { |name| trait_by_name(name) }
+    rescue KeyError => error
+      raise error_with_definition_name(error)
+    end
+
+    def error_with_definition_name(error)
+      message = error.message
+      message.insert(
+        message.index("\nDid you mean?") || message.length,
+        " referenced within \"#{name}\" definition"
+      )
+
+      error.class.new(message).tap do |new_error|
+        new_error.set_backtrace(error.backtrace)
+      end
     end
 
     def additional_traits
diff --git a/spec/acceptance/traits_spec.rb b/spec/acceptance/traits_spec.rb
index 3f1cfba..293f963 100644
--- a/spec/acceptance/traits_spec.rb
+++ b/spec/acceptance/traits_spec.rb
@@ -9,10 +9,6 @@ describe "an instance generated by a factory with multiple traits" do
       great: :string)
 
     FactoryBot.define do
-      factory :user_without_admin_scoping, class: User do
-        admin_trait
-      end
-
       factory :user do
         name { "John" }
 
@@ -187,15 +183,6 @@ describe "an instance generated by a factory with multiple traits" do
     its(:date_of_birth) { should eq Date.parse("1/1/2000") }
   end
 
-  context "factory outside of scope" do
-    subject { FactoryBot.create(:user_without_admin_scoping) }
-
-    it "raises an error" do
-      expect { subject }
-        .to raise_error(KeyError, "Trait not registered: \"admin_trait\"")
-    end
-  end
-
   context "child factory using grandparents' trait" do
     subject { FactoryBot.create(:female_great_user) }
     its(:great) { should eq "GREAT!!!" }
@@ -293,15 +280,81 @@ describe "trait indifferent access" do
 end
 
 describe "looking up traits that don't exist" do
-  it "raises a KeyError" do
-    define_class("User")
+  context "when passing an invalid override trait" do
+    it "raises a KeyError" do
+      define_class("User")
 
-    FactoryBot.define do
-      factory :user
+      FactoryBot.define do
+        factory :user
+      end
+
+      expect { FactoryBot.build(:user, double("not a trait")) }
+        .to raise_error(KeyError)
+    end
+  end
+
+  context "when the factory includes an invalid default trait" do
+    it "raises a KeyError including the factory name" do
+      define_class("User")
+
+      FactoryBot.define do
+        factory :user do
+          inaccessible_trait
+        end
+
+        factory :some_other_factory do
+          trait :inaccessible_trait
+        end
+      end
+
+      expect { FactoryBot.build(:user) }.to raise_error(
+        KeyError,
+        'Trait not registered: "inaccessible_trait" referenced within "user" definition'
+      )
     end
 
-    expect { FactoryBot.build(:user, double("not a trait")) }
-      .to raise_error(KeyError)
+    it "maintains 'Did you mean?' suggestions at the end of the error message" do
+      define_class("User")
+
+      FactoryBot.define do
+        trait :not_quit
+
+        factory :user do
+          not_quite
+        end
+      end
+
+      expect { FactoryBot.build(:user) }.to raise_error(
+        KeyError,
+        <<~MSG.strip
+          Trait not registered: "not_quite" referenced within "user" definition
+          Did you mean?  "not_quit"
+        MSG
+      )
+    end
+  end
+
+  context "when a trait includes an invalid default trait" do
+    it "raises a KeyError including the factory name" do
+      define_class("User")
+
+      FactoryBot.define do
+        factory :user do
+          trait :admin do
+            inaccessible_trait
+          end
+        end
+
+        factory :some_other_factory do
+          trait :inaccessible_trait
+        end
+      end
+
+      expect { FactoryBot.build(:user, :admin) }.to raise_error(
+        KeyError,
+        'Trait not registered: "inaccessible_trait" referenced within "admin" definition'
+      )
+    end
   end
 end