Merge pull request #1871 from TobiasBales/better-fixtures-types

Infer types for active record fixtures

lib/tapioca/dsl/compilers/active_record_fixtures.rb

@@ -1,7 +1,9 @@
 # typed: strict
 # frozen_string_literal: true
 
-return unless defined?(Rails) && defined?(ActiveSupport::TestCase) && defined?(ActiveRecord::TestFixtures)
+return unless defined?(Rails) &&
+  defined?(ActiveSupport::TestCase) &&
+  defined?(ActiveRecord::TestFixtures)
 
 module Tapioca
   module Dsl
@@ -24,8 +26,10 @@ module Compilers
       # # test_case.rbi
       # # typed: true
       # class ActiveSupport::TestCase
-      #   sig { params(fixture_names: T.any(String, Symbol)).returns(T.untyped) }
-      #   def posts(*fixture_names); end
+      #   sig { params(fixture_name: T.any(String, Symbol), other_fixtures: NilClass).returns(Post) }
+      #   sig { params(fixture_name: T.any(String, Symbol), other_fixtures: T.any(String, Symbol))
+      #           .returns(T::Array[Post]) }
+      #   def posts(fixture_name, *other_fixtures); end
       # end
       # ~~~
       class ActiveRecordFixtures < Compiler
@@ -104,10 +108,90 @@ def method_names_from_eager_fixture_loader
 
         sig { params(mod: RBI::Scope, name: String).void }
         def create_fixture_method(mod, name)
-          mod.create_method(
-            name,
-            parameters: [create_rest_param("fixture_names", type: "T.any(String, Symbol)")],
-            return_type: "T.untyped",
+          return_type = return_type_for_fixture(name)
+          mod << RBI::Method.new(name) do |node|
+            node.add_param("fixture_name")
+            node.add_rest_param("other_fixtures")
+
+            node.add_sig do |sig|
+              sig.add_param("fixture_name", "T.any(String, Symbol)")
+              sig.add_param("other_fixtures", "NilClass")
+              sig.return_type = return_type
+            end
+
+            node.add_sig do |sig|
+              sig.add_param("fixture_name", "T.any(String, Symbol)")
+              sig.add_param("other_fixtures", "T.any(String, Symbol)")
+              sig.return_type = "T::Array[#{return_type}]"
+            end
+          end
+        end
+
+        sig { params(fixture_name: String).returns(String) }
+        def return_type_for_fixture(fixture_name)
+          fixture_class_mapping_from_fixture_files[fixture_name] ||
+            fixture_class_from_fixture_set(fixture_name) ||
+            fixture_class_from_active_record_base_class_mapping[fixture_name] ||
+            "T.untyped"
+        end
+
+        sig { params(fixture_name: String).returns(T.nilable(String)) }
+        def fixture_class_from_fixture_set(fixture_name)
+          # only rails 7.1+ support fixture sets so this is conditional
+          return unless fixture_loader.respond_to?(:fixture_sets)
+
+          model_name_from_fixture_set = T.unsafe(fixture_loader).fixture_sets[fixture_name]
+          return unless model_name_from_fixture_set
+
+          model_name = ActiveRecord::FixtureSet.default_fixture_model_name(model_name_from_fixture_set)
+          return unless Object.const_defined?(model_name)
+
+          model_name
+        end
+
+        sig { returns(T::Hash[String, String]) }
+        def fixture_class_from_active_record_base_class_mapping
+          @fixture_class_mapping ||= T.let(
+            begin
+              ActiveRecord::Base.descendants.each_with_object({}) do |model_class, mapping|
+                class_name = model_class.name
+
+                fixture_name = class_name.underscore.gsub("/", "_")
+                fixture_name = fixture_name.pluralize if ActiveRecord::Base.pluralize_table_names
+
+                mapping[fixture_name] = class_name
+
+                mapping
+              end
+            end,
+            T.nilable(T::Hash[String, String]),
+          )
+        end
+
+        sig { returns(T::Hash[String, String]) }
+        def fixture_class_mapping_from_fixture_files
+          @fixture_file_class_mapping ||= T.let(
+            begin
+              fixture_paths = if T.unsafe(fixture_loader).respond_to?(:fixture_paths)
+                T.unsafe(fixture_loader).fixture_paths
+              else
+                T.unsafe(fixture_loader).fixture_path
+              end
+
+              Array(fixture_paths).each_with_object({}) do |path, mapping|
+                Dir["#{path}{.yml,/{**,*}/*.yml}"].select do |file|
+                  next unless ::File.file?(file)
+
+                  ActiveRecord::FixtureSet::File.open(file) do |fh|
+                    next unless fh.model_class
+
+                    fixture_name = file.delete_prefix(path.to_s).delete_prefix("/").delete_suffix(".yml")
+                    mapping[fixture_name] = fh.model_class
+                  end
+                end
+              end
+            end,
+            T.nilable(T::Hash[String, String]),
           )
         end
       end

manual/compiler_activerecordfixtures.md

@@ -18,7 +18,9 @@ The generated RBI by this compiler will produce the following
 # test_case.rbi
 # typed: true
 class ActiveSupport::TestCase
-  sig { params(fixture_names: T.any(String, Symbol)).returns(T.untyped) }
-  def posts(*fixture_names); end
+  sig { params(fixture_name: T.any(String, Symbol), other_fixtures: NilClass).returns(Post) }
+  sig { params(fixture_name: T.any(String, Symbol), other_fixtures: T.any(String, Symbol))
+          .returns(T::Array[Post]) }
+  def posts(fixture_name, *other_fixtures); end
 end
 ~~~

sorbet/rbi/shims/active_record_fixture_set.rbi

@@ -0,0 +1,16 @@
+# typed: strict
+
+# ActiveRecord::TestFixtures can't be loaded outside of a Rails application
+
+class ActiveRecord::FixtureSet
+  sig { params(name: String).returns(String) }
+  def self.default_fixture_model_name(name); end
+end
+
+class ActiveRecord::FixtureSet::File
+  sig { returns(T.nilable(String)) }
+  def modle_class; end
+
+  sig params(filename: String, blk: T.proc.params(arg0: ActiveRecord::FixtureSet::File).void)
+  def self.open(filename, &block); end
+end

spec/tapioca/dsl/compilers/active_record_fixtures_spec.rb

@@ -64,20 +64,34 @@ class User
                   updated_at: 2021-09-08 11:00:00
               YAML
 
+              add_ruby_file("test_models.rb", <<~RUBY)
+                class Post < ActiveRecord::Base
+                end
+              RUBY
+
               expected = <<~RBI
                 # typed: strong
 
                 class ActiveSupport::TestCase
-                  sig { params(fixture_names: T.any(String, Symbol)).returns(T.untyped) }
-                  def posts(*fixture_names); end
+                  sig { params(fixture_name: T.any(String, Symbol), other_fixtures: NilClass).returns(Post) }
+                  sig { params(fixture_name: T.any(String, Symbol), other_fixtures: T.any(String, Symbol)).returns(T::Array[Post]) }
+                  def posts(fixture_name, *other_fixtures); end
                 end
               RBI
 
               assert_equal(expected, rbi_for("ActiveSupport::TestCase"))
             end
 
             it "generates methods for fixtures from multiple sources" do
-              add_content_file("test/fixtures/posts.yml", <<~YAML)
+              add_ruby_file("test_models.rb", <<~RUBY)
+                module Blog
+                  class Post < ActiveRecord::Base
+                  end
+                end
+                class User < ActiveRecord::Base
+                end
+              RUBY
+              add_content_file("test/fixtures/blog/posts.yml", <<~YAML)
                 super_post:
                   title: An incredible Ruby post
                   author: Johnny Developer
@@ -97,11 +111,64 @@ def posts(*fixture_names); end
                 # typed: strong
 
                 class ActiveSupport::TestCase
-                  sig { params(fixture_names: T.any(String, Symbol)).returns(T.untyped) }
-                  def posts(*fixture_names); end
+                  sig { params(fixture_name: T.any(String, Symbol), other_fixtures: NilClass).returns(Blog::Post) }
+                  sig { params(fixture_name: T.any(String, Symbol), other_fixtures: T.any(String, Symbol)).returns(T::Array[Blog::Post]) }
+                  def blog_posts(fixture_name, *other_fixtures); end
+
+                  sig { params(fixture_name: T.any(String, Symbol), other_fixtures: NilClass).returns(User) }
+                  sig { params(fixture_name: T.any(String, Symbol), other_fixtures: T.any(String, Symbol)).returns(T::Array[User]) }
+                  def users(fixture_name, *other_fixtures); end
+                end
+              RBI
+
+              assert_equal(expected, rbi_for("ActiveSupport::TestCase"))
+            end
+
+            it "generates methods for fixtures with explicit class name" do
+              add_content_file("test/fixtures/posts_with_other_names.yml", <<~YAML)
+                _fixture:
+                  model_class: Post
+                super_post:
+                  title: An incredible Ruby post
+                  author: Johnny Developer
+                  created_at: 2021-09-08 11:00:00
+                  updated_at: 2021-09-08 11:00:00
+              YAML
+
+              add_ruby_file("test_models.rb", <<~RUBY)
+                class Post < ActiveRecord::Base
+                end
+              RUBY
 
-                  sig { params(fixture_names: T.any(String, Symbol)).returns(T.untyped) }
-                  def users(*fixture_names); end
+              expected = <<~RBI
+                # typed: strong
+
+                class ActiveSupport::TestCase
+                  sig { params(fixture_name: T.any(String, Symbol), other_fixtures: NilClass).returns(Post) }
+                  sig { params(fixture_name: T.any(String, Symbol), other_fixtures: T.any(String, Symbol)).returns(T::Array[Post]) }
+                  def posts_with_other_names(fixture_name, *other_fixtures); end
+                end
+              RBI
+
+              assert_equal(expected, rbi_for("ActiveSupport::TestCase"))
+            end
+
+            it "generates methods for fixtures with a fallback to T.untyped if no matching model exists" do
+              add_content_file("test/fixtures/posts.yml", <<~YAML)
+                super_post:
+                  title: An incredible Ruby post
+                  author: Johnny Developer
+                  created_at: 2021-09-08 11:00:00
+                  updated_at: 2021-09-08 11:00:00
+              YAML
+
+              expected = <<~RBI
+                # typed: strong
+
+                class ActiveSupport::TestCase
+                  sig { params(fixture_name: T.any(String, Symbol), other_fixtures: NilClass).returns(T.untyped) }
+                  sig { params(fixture_name: T.any(String, Symbol), other_fixtures: T.any(String, Symbol)).returns(T::Array[T.untyped]) }
+                  def posts(fixture_name, *other_fixtures); end
                 end
               RBI