Added STI support to init and building associations

Allows you to do BaseClass.new(:type => "SubClass") as well as
parent.children.build(:type => "SubClass") or parent.build_child
to initialize an STI subclass. Ensures that the class name is a
valid class and that it is in the ancestors of the super class
that the association is expecting.

7 files changed, 126 insertions, 0 deletions

diff --git a/activerecord/test/cases/associations/belongs_to_associations_test.rb b/activerecord/test/cases/associations/belongs_to_associations_test.rb
index 5f7825783b..49d6c31c9a 100644
--- a/activerecord/test/cases/associations/belongs_to_associations_test.rb
+++ b/activerecord/test/cases/associations/belongs_to_associations_test.rb
@@ -109,6 +109,34 @@ class BelongsToAssociationsTest < ActiveRecord::TestCase
     assert_equal apple.id, citibank.firm_id
   end
 
+  def test_building_the_belonging_object_with_implicit_sti_base_class
+    account = Account.new
+    company = account.build_firm
+    assert(company.kind_of?(Company), "Expected #{company.class} to be a Company")
+  end
+
+  def test_building_the_belonging_object_with_explicit_sti_base_class
+    account = Account.new
+    company = account.build_firm(:type => "Company")
+    assert(company.kind_of?(Company), "Expected #{company.class} to be a Company")
+  end
+
+  def test_building_the_belonging_object_with_sti_subclass
+    account = Account.new
+    company = account.build_firm(:type => "Firm")
+    assert(company.kind_of?(Firm), "Expected #{company.class} to be a Firm")
+  end
+
+  def test_building_the_belonging_object_with_an_invalid_type
+    account = Account.new
+    assert_raise(ActiveRecord::SubclassNotFound) { account.build_firm(:type => "InvalidType") }
+  end
+
+  def test_building_the_belonging_object_with_an_unrelated_type
+    account = Account.new
+    assert_raise(ActiveRecord::SubclassNotFound) { account.build_firm(:type => "Account") }
+  end
+
   def test_building_the_belonging_object_with_primary_key
     client = Client.create(:name => "Primary key client")
     apple  = client.build_firm_with_primary_key("name" => "Apple")
diff --git a/activerecord/test/cases/associations/has_many_associations_test.rb b/activerecord/test/cases/associations/has_many_associations_test.rb
index e5022d49f1..eee0cf03aa 100644
--- a/activerecord/test/cases/associations/has_many_associations_test.rb
+++ b/activerecord/test/cases/associations/has_many_associations_test.rb
@@ -144,6 +144,34 @@ class HasManyAssociationsTest < ActiveRecord::TestCase
     assert_equal 'defaulty', bulb.name
   end
 
+  def test_building_the_associated_object_with_implicit_sti_base_class
+    firm = DependentFirm.new
+    company = firm.companies.build
+    assert(company.kind_of?(Company), "Expected #{company.class} to be a Company")
+  end
+
+  def test_building_the_associated_object_with_explicit_sti_base_class
+    firm = DependentFirm.new
+    company = firm.companies.build(:type => "Company")
+    assert(company.kind_of?(Company), "Expected #{company.class} to be a Company")
+  end
+
+  def test_building_the_associated_object_with_sti_subclass
+    firm = DependentFirm.new
+    company = firm.companies.build(:type => "Client")
+    assert(company.kind_of?(Client), "Expected #{company.class} to be a Client")
+  end
+
+  def test_building_the_associated_object_with_an_invalid_type
+    firm = DependentFirm.new
+    assert_raise(ActiveRecord::SubclassNotFound) { firm.companies.build(:type => "Invalid") }
+  end
+
+  def test_building_the_associated_object_with_an_unrelated_type
+    firm = DependentFirm.new
+    assert_raise(ActiveRecord::SubclassNotFound) { firm.companies.build(:type => "Account") }
+  end
+
   def test_association_keys_bypass_attribute_protection
     car = Car.create(:name => 'honda')
 
diff --git a/activerecord/test/cases/associations/has_one_associations_test.rb b/activerecord/test/cases/associations/has_one_associations_test.rb
index ea1cfa0805..5c4f6bbb32 100644
--- a/activerecord/test/cases/associations/has_one_associations_test.rb
+++ b/activerecord/test/cases/associations/has_one_associations_test.rb
@@ -6,6 +6,8 @@ require 'models/ship'
 require 'models/pirate'
 require 'models/car'
 require 'models/bulb'
+require 'models/author'
+require 'models/post'
 
 class HasOneAssociationsTest < ActiveRecord::TestCase
   self.use_transactional_fixtures = false unless supports_savepoints?
@@ -212,6 +214,34 @@ class HasOneAssociationsTest < ActiveRecord::TestCase
     }
   end
 
+  def test_building_the_associated_object_with_implicit_sti_base_class
+    firm = DependentFirm.new
+    company = firm.build_company
+    assert(company.kind_of?(Company), "Expected #{company.class} to be a Company")
+  end
+
+  def test_building_the_associated_object_with_explicit_sti_base_class
+    firm = DependentFirm.new
+    company = firm.build_company(:type => "Company")
+    assert(company.kind_of?(Company), "Expected #{company.class} to be a Company")
+  end
+
+  def test_building_the_associated_object_with_sti_subclass
+    firm = DependentFirm.new
+    company = firm.build_company(:type => "Client")
+    assert(company.kind_of?(Client), "Expected #{company.class} to be a Client")
+  end
+
+  def test_building_the_associated_object_with_an_invalid_type
+    firm = DependentFirm.new
+    assert_raise(ActiveRecord::SubclassNotFound) { firm.build_company(:type => "Invalid") }
+  end
+
+  def test_building_the_associated_object_with_an_unrelated_type
+    firm = DependentFirm.new
+    assert_raise(ActiveRecord::SubclassNotFound) { firm.build_company(:type => "Account") }
+  end
+
   def test_build_and_create_should_not_happen_within_scope
     pirate = pirates(:blackbeard)
     scoped_count = pirate.association(:foo_bulb).scope.where_values.count
diff --git a/activerecord/test/cases/forbidden_attributes_protection_test.rb b/activerecord/test/cases/forbidden_attributes_protection_test.rb
index 9a2172f41e..490b599fb6 100644
--- a/activerecord/test/cases/forbidden_attributes_protection_test.rb
+++ b/activerecord/test/cases/forbidden_attributes_protection_test.rb
@@ -1,6 +1,7 @@
 require 'cases/helper'
 require 'active_support/core_ext/hash/indifferent_access'
 require 'models/person'
+require 'models/company'
 
 class ProtectedParams < ActiveSupport::HashWithIndifferentAccess
   attr_accessor :permitted
@@ -40,6 +41,20 @@ class ForbiddenAttributesProtectionTest < ActiveRecord::TestCase
     assert_equal 'm', person.gender
   end
 
+  def test_forbidden_attributes_cannot_be_used_for_sti_inheritance_column
+    params = ProtectedParams.new(type: 'Client')
+    assert_raises(ActiveModel::ForbiddenAttributesError) do
+      Company.new(params)
+    end
+  end
+
+  def test_permitted_attributes_can_be_used_for_sti_inheritance_column
+    params = ProtectedParams.new(type: 'Client')
+    params.permit!
+    person = Company.new(params)
+    assert_equal person.class, Client
+  end
+
   def test_regular_hash_should_still_be_used_for_mass_assignment
     person = Person.new(first_name: 'Guille', gender: 'm')
 
diff --git a/activerecord/test/cases/inheritance_test.rb b/activerecord/test/cases/inheritance_test.rb
index aab7aa51dd..2466a764f6 100644
--- a/activerecord/test/cases/inheritance_test.rb
+++ b/activerecord/test/cases/inheritance_test.rb
@@ -156,6 +156,29 @@ class InheritanceTest < ActiveRecord::TestCase
     assert_kind_of Cabbage, savoy
   end
 
+  def test_inheritance_new_with_default_class
+    company = Company.new
+    assert_equal company.class, Company
+  end
+
+  def test_inheritance_new_with_base_class
+    company = Company.new(:type => 'Company')
+    assert_equal company.class, Company
+  end
+
+  def test_inheritance_new_with_subclass
+    firm = Company.new(:type => 'Firm')
+    assert_equal firm.class, Firm
+  end
+
+  def test_new_with_invalid_type
+    assert_raise(ActiveRecord::SubclassNotFound) { Company.new(:type => 'InvalidType') }
+  end
+
+  def test_new_with_unrelated_type
+    assert_raise(ActiveRecord::SubclassNotFound) { Company.new(:type => 'Account') }
+  end
+
   def test_inheritance_condition
     assert_equal 10, Company.count
     assert_equal 2, Firm.count
diff --git a/activerecord/test/models/author.rb b/activerecord/test/models/author.rb
index 77f4a2ec87..6935cfb0ea 100644
--- a/activerecord/test/models/author.rb
+++ b/activerecord/test/models/author.rb
@@ -1,5 +1,6 @@
 class Author < ActiveRecord::Base
   has_many :posts
+  has_one :post
   has_many :very_special_comments, :through => :posts
   has_many :posts_with_comments, -> { includes(:comments) }, :class_name => "Post"
   has_many :popular_grouped_posts, -> { includes(:comments).group("type").having("SUM(comments_count) > 1").select("type") }, :class_name => "Post"
diff --git a/activerecord/test/models/company.rb b/activerecord/test/models/company.rb
index 17b17724e8..3ca8f69646 100644
--- a/activerecord/test/models/company.rb
+++ b/activerecord/test/models/company.rb
@@ -111,6 +111,7 @@ end
 class DependentFirm < Company
   has_one :account, :foreign_key => "firm_id", :dependent => :nullify
   has_many :companies, :foreign_key => 'client_of', :dependent => :nullify
+  has_one :company, :foreign_key => 'client_of', :dependent => :nullify
 end
 
 class RestrictedFirm < Company