Add tests for `TypeMap#fetch` and push up to `TypeMap`

It doesn't make sense for the subclass to implement this method, and not
have it on the parent. We can also DRY up the implementation of
`#lookup` to be defined in terms of fetch, which will give us a single
point of entry

3 files changed, 20 insertions, 5 deletions

diff --git a/activerecord/lib/active_record/type/hash_lookup_type_map.rb b/activerecord/lib/active_record/type/hash_lookup_type_map.rb
index bf92680268..90a8c674cb 100644
--- a/activerecord/lib/active_record/type/hash_lookup_type_map.rb
+++ b/activerecord/lib/active_record/type/hash_lookup_type_map.rb
@@ -3,10 +3,6 @@ module ActiveRecord
     class HashLookupTypeMap < TypeMap # :nodoc:
       delegate :key?, to: :@mapping
 
-      def lookup(type, *args)
-        @mapping.fetch(type, proc { default_value }).call(type, *args)
-      end
-
       def fetch(type, *args, &block)
         @mapping.fetch(type, block).call(type, *args)
       end
diff --git a/activerecord/lib/active_record/type/type_map.rb b/activerecord/lib/active_record/type/type_map.rb
index 88c5f9c497..13d4943861 100644
--- a/activerecord/lib/active_record/type/type_map.rb
+++ b/activerecord/lib/active_record/type/type_map.rb
@@ -6,6 +6,10 @@ module ActiveRecord
       end
 
       def lookup(lookup_key, *args)
+        fetch(lookup_key, *args) { default_value }
+      end
+
+      def fetch(lookup_key, *args)
         matching_pair = @mapping.reverse_each.detect do |key, _|
           key === lookup_key
         end
@@ -13,7 +17,7 @@ module ActiveRecord
         if matching_pair
           matching_pair.last.call(lookup_key, *args)
         else
-          default_value
+          yield lookup_key, *args
         end
       end
 
diff --git a/activerecord/test/cases/type/type_map_test.rb b/activerecord/test/cases/type/type_map_test.rb
index 4e32f92dd0..1e4eba1961 100644
--- a/activerecord/test/cases/type/type_map_test.rb
+++ b/activerecord/test/cases/type/type_map_test.rb
@@ -124,6 +124,21 @@ module ActiveRecord
         assert_equal mapping.lookup(3), 'string'
         assert_kind_of Type::Value, mapping.lookup(4)
       end
+
+      def test_fetch
+        mapping = TypeMap.new
+        mapping.register_type(1, "string")
+
+        assert_equal "string", mapping.fetch(1) { "int" }
+        assert_equal "int", mapping.fetch(2) { "int" }
+      end
+
+      def test_fetch_yields_args
+        mapping = TypeMap.new
+
+        assert_equal "foo-1-2-3", mapping.fetch("foo", 1, 2, 3) { |*args| args.join("-") }
+        assert_equal "bar-1-2-3", mapping.fetch("bar", 1, 2, 3) { |*args| args.join("-") }
+      end
     end
   end
 end