Merge pull request rails#47947 from Shopify/pm/fixtures-with-composite-key

Support building fixtures with associations that use CPK

Author: eileencodes

activerecord/lib/active_record/fixture_set/table_row.rb

@@ -125,13 +125,11 @@ def generate_primary_key
         end
 
         def generate_composite_primary_key
-          id = ActiveRecord::FixtureSet.identify(@label)
-          model_metadata.primary_key_name.each_with_index do |column, index|
+          composite_key = ActiveRecord::FixtureSet.composite_identify(@label, model_metadata.primary_key_name)
+          composite_key.each do |column, value|
             next if column_defined?(column)
-            raise "Automatic key generation assumes columns of type Integer." unless model_metadata.column_type(column) == :integer
 
-            # Shift label identifier index-#-of-times to differentiate sub-components in deterministic manner.
-            @row[column] = (id << index) % ActiveRecord::FixtureSet::MAX_ID
+            @row[column] = value
           end
         end
 
@@ -165,8 +163,17 @@ def resolve_sti_reflections
                   raise PrimaryKeyError.new(@label, association, value)
                 end
 
-                fk_type = reflection_class.type_for_attribute(fk_name).type
-                @row[fk_name] = ActiveRecord::FixtureSet.identify(value, fk_type)
+                if fk_name.is_a?(Array)
+                  composite_key = ActiveRecord::FixtureSet.composite_identify(value, fk_name)
+                  composite_key.each do |column, value|
+                    next if column_defined?(column)
+
+                    @row[column] = value
+                  end
+                else
+                  fk_type = reflection_class.type_for_attribute(fk_name).type
+                  @row[fk_name] = ActiveRecord::FixtureSet.identify(value, fk_type)
+                end
               end
             when :has_many
               if association.options[:through]

activerecord/lib/active_record/fixtures.rb

@@ -589,6 +589,18 @@ def identify(label, column_type = :integer)
         end
       end
 
+      # Returns a consistent, platform-independent hash representing a mapping
+      # between the label and the subcomponents of the provided composite key.
+      #
+      # Example:
+      # composite_identify("label", [:a, :b, :c]) => { a: hash_1, b: hash_2, c: hash_3 }
+      def composite_identify(label, key)
+        key
+          .index_with
+          .with_index { |sub_key, index| (identify(label) << index) % MAX_ID }
+          .with_indifferent_access
+      end
+
       # Superclass for the evaluation contexts used by ERB fixtures.
       def context_class
         @context_class ||= Class.new

activerecord/test/cases/fixtures_test.rb

@@ -1643,15 +1643,15 @@ def readonly_config
   end
 
   class CompositePkFixturesTest < ActiveRecord::TestCase
-    fixtures :cpk_orders, :cpk_books, :authors
+    fixtures :cpk_orders, :cpk_books, :cpk_authors, :cpk_reviews, :cpk_order_agreements
 
     def test_generates_composite_primary_key_for_partially_filled_fixtures
-      david = authors(:david)
-      david_cpk_book = cpk_books(:cpk_known_author_david_book)
+      alice = cpk_authors(:cpk_great_author)
+      alice_cpk_book = cpk_books(:cpk_great_author_first_book)
 
-      assert_not_empty(david_cpk_book.id.compact)
-      assert_equal david.id, david_cpk_book.author_id
-      assert_not_nil david_cpk_book.number
+      assert_not_empty(alice_cpk_book.id.compact)
+      assert_equal alice.id, alice_cpk_book.author_id
+      assert_not_nil alice_cpk_book.number
     end
 
     def test_generates_composite_primary_key_ids
@@ -1664,5 +1664,70 @@ def test_generates_composite_primary_key_ids
     def test_generates_composite_primary_key_with_unique_components
       assert_equal 2, cpk_orders(:cpk_groceries_order_1).id.uniq.size
     end
+
+    def test_resolves_associations_using_composite_primary_keys
+      review = cpk_reviews(:first_book_review)
+      generated_book = cpk_books(:cpk_book_with_generated_pk)
+
+      assert_equal generated_book.id, [review.author_id, review.number]
+      assert_equal generated_book, review.book
+    end
+
+    def test_resolves_associations_using_composite_primary_keys_with_partially_filled_values
+      review = cpk_reviews(:second_book_review_for_book_with_partial_pk_defined)
+      book_with_partially_filled_cpk = cpk_books(:cpk_great_author_first_book)
+
+      assert_equal book_with_partially_filled_cpk.id, [review.author_id, review.number]
+      assert_equal book_with_partially_filled_cpk, review.book
+    end
+
+    def test_association_with_custom_primary_key
+      order = cpk_orders(:cpk_groceries_order_2)
+      order_agreement = cpk_order_agreements(:order_agreement_three)
+
+      _, order_id = order.id
+
+      assert_equal order_id, order_agreement.order_id
+      assert_equal order, order_agreement.order
+    end
+
+    def test_composite_identify_resolves_to_same_values
+      identify_one = ActiveRecord::FixtureSet.composite_identify("label", [:a, :b, :c])
+      identify_two = ActiveRecord::FixtureSet.composite_identify("label", [:a, :b, :c])
+
+      assert_equal identify_one, identify_two
+    end
+
+    def test_composite_identify_returns_hash_with_key_names
+      id = ActiveRecord::FixtureSet.composite_identify("order", Cpk::Order.primary_key)
+
+      assert_equal ["shop_id", "id"], id.keys
+    end
+
+    def test_composite_identify_uses_same_hashing_algorithm_as_identify_for_first_attribute
+      id_hash = ActiveRecord::FixtureSet.composite_identify("order", [:first_attribute, :second_attribute])
+      id = ActiveRecord::FixtureSet.identify("order")
+
+      assert_equal id, id_hash[:first_attribute]
+      assert_not_equal id, id_hash[:second_attribute]
+    end
+
+    def test_composite_identify_hashes_one_label_to_same_values_irrespective_of_column_names
+      id_hash_one = ActiveRecord::FixtureSet.composite_identify("order", [:first_attribute, :second_attribute])
+      id_hash_two = ActiveRecord::FixtureSet.composite_identify("order", [:shop_id, :id])
+
+      assert_equal id_hash_one.values, id_hash_two.values
+      assert_not_equal id_hash_one.keys, id_hash_two.keys
+    end
+
+    def test_composite_identify_hashes_to_same_values_based_on_position_in_key
+      id = ActiveRecord::FixtureSet.identify("order")
+      id_hash_two = ActiveRecord::FixtureSet.composite_identify("order", [:one, :two])
+      id_hash_three = ActiveRecord::FixtureSet.composite_identify("order", [:one, :two, :three])
+
+      assert_equal id, id_hash_two.values.first
+      assert_equal id, id_hash_three.values.first
+      assert_equal id_has_two.values, id_hash_three.values.slice(0, 2)
+    end
   end
 end

activerecord/test/fixtures/cpk_books.yml

@@ -16,7 +16,6 @@ cpk_famous_author_first_book:
   title: "Ruby on Rails"
   revision: 1
 
-cpk_known_author_david_book:
-  author_id: 1
-  title: "David's CPK Book"
+cpk_book_with_generated_pk:
+  title: "Generated author's book"
   revision: 1

activerecord/test/fixtures/cpk_order_agreements.yml

@@ -6,3 +6,8 @@ order_agreement_one:
 
 order_agreement_two:
   signature: "xyz789"
+
+order_agreement_three:
+  order_id: <%= ActiveRecord::FixtureSet.composite_identify(:cpk_groceries_order_2, Cpk::Order.primary_key)[:id] %>
+  signature: "def321"
+

activerecord/test/fixtures/cpk_reviews.yml

@@ -0,0 +1,13 @@
+_fixture:
+  model_class: Cpk::Review
+
+first_book_review:
+  book: cpk_book_with_generated_pk
+  rating: 5
+  comment: "The first book was alright."
+
+second_book_review_for_book_with_partial_pk_defined:
+  book: cpk_great_author_first_book
+  author_id: <%= ActiveRecord::FixtureSet.identify(:cpk_great_author) %>
+  rating: 5
+  comment: "The first book was alright."

activerecord/test/models/cpk.rb

@@ -3,4 +3,5 @@
 require_relative "cpk/author"
 require_relative "cpk/book"
 require_relative "cpk/order"
+require_relative "cpk/review"
 require_relative "cpk/order_agreement"

activerecord/test/models/cpk/review.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module Cpk
+  class Review < ActiveRecord::Base
+    self.table_name = :cpk_reviews
+
+    belongs_to :book, class_name: "Cpk::Book", query_constraints: [:author_id, :number]
+  end
+end

activerecord/test/schema/schema.rb

@@ -251,6 +251,13 @@
     t.string :name
   end
 
+  create_table :cpk_reviews, force: true do |t|
+    t.integer :author_id
+    t.integer :number
+    t.integer :rating
+    t.string :comment
+  end
+
   create_table :cpk_orders, primary_key: [:shop_id, :id], force: true do |t|
     t.integer :shop_id
     t.integer :id