sorbet-rails 0.6.3 → 0.7.1

data/lib/sorbet-rails/tasks/rails_rbi.rake

@@ -35,19 +35,19 @@ namespace :rails_rbi do
     File.write(file_path, inspector.format(SorbetRails::RoutesRbiFormatter.new))
   end
 
-  desc "Copy custom rbis for ActionController::Parameters, pluck_to_struct, etc."
+  desc "Copy custom rbis for typed_params, pluck_to_struct, etc."
   task custom: :environment do
     copy_bundled_rbi('type_assert.rbi')
-    copy_bundled_rbi('parameters.rbi')
     copy_bundled_rbi('pluck_to_tstruct.rbi')
     copy_bundled_rbi('typed_params.rbi')
+    copy_bundled_rbi('typed_enum.rbi')
 
     # These files were previously bundled_rbi but are now generated so this
     # is needed for backwards compatibility with anyone using `rails_rbi:custom`
     Rake::Task['rails_rbi:active_record'].invoke
   end
 
-  desc "Generate rbis for rails mailers"
+  desc "Generate rbis for rails active_record base"
   task :active_record, [:root_dir] => :environment do |t, args|
     formatter = SorbetRails::ActiveRecordRbiFormatter.new
     FileUtils.mkdir_p(Rails.root.join("sorbet", "rails-rbi"))
@@ -77,11 +77,19 @@ namespace :rails_rbi do
       models_to_generate = models_to_generate - blacklisted_models
     end
 
-    generated_rbis = generate_rbis_for_models(models_to_generate, all_models)
-    generated_rbis.each do |model_name, contents|
-      file_path = Rails.root.join("sorbet", "rails-rbi", "models", "#{model_name.underscore}.rbi")
-      FileUtils.mkdir_p(File.dirname(file_path))
-      File.write(file_path, contents)
+    available_class_names = Set.new(all_models.map { |c| c.name })
+    models_to_generate.each do |model_class|
+      model_class_name = model_class.to_s
+      begin
+        formatter = SorbetRails::ModelRbiFormatter.new(model_class, available_class_names)
+        file_path = Rails.root.join("sorbet", "rails-rbi", "models", "#{model_class_name.underscore}.rbi")
+        FileUtils.mkdir_p(File.dirname(file_path))
+        File.write(file_path, formatter.generate_rbi)
+      rescue StandardError, NotImplementedError => ex
+        puts "---"
+        puts "Error when handling model #{model_class_name}: #{ex}"
+        nil
+      end
     end
   end
 
@@ -121,7 +129,7 @@ namespace :rails_rbi do
         "mailers",
         "#{mailer_class.name.underscore}.rbi",
       )
-      formatter = SorbetRails::MailerRbiFormatter.new(mailer_class)
+      formatter = ::SorbetRails.config.mailer_generator_class.new(mailer_class)
       FileUtils.mkdir_p(File.dirname(file_path))
       File.write(file_path, formatter.generate_rbi)
     end
@@ -139,28 +147,12 @@ namespace :rails_rbi do
         "jobs",
         "#{job_class.name.underscore}.rbi",
       )
-      formatter = SorbetRails::JobRbiFormatter.new(job_class)
+      formatter = ::SorbetRails.config.job_generator_class.new(job_class)
       FileUtils.mkdir_p(File.dirname(file_path))
       File.write(file_path, formatter.generate_rbi)
     end
   end
 
-  def generate_rbis_for_models(model_classes, available_classes)
-    available_class_names = Set.new(available_classes.map { |c| c.name })
-    formatted = model_classes.map do |model_class|
-      model_class_name = model_class.to_s
-      begin
-        formatter = SorbetRails::ModelRbiFormatter.new(model_class, available_class_names)
-        [model_class_name, formatter.generate_rbi]
-      rescue StandardError, NotImplementedError => ex
-        puts "---"
-        puts "Error when handling model #{model_class_name}: #{ex}"
-        nil
-      end
-    end
-    Hash[formatted.compact] # remove models with errors
-  end
-
   def blacklisted_models
     blacklisted_models = []
     blacklisted_models << ApplicationRecord if defined?(ApplicationRecord)

data/sorbet-rails.gemspec

@@ -1,6 +1,6 @@
 Gem::Specification.new do |s|
   s.name          = %q{sorbet-rails}
-  s.version       = "0.6.3"
+  s.version       = "0.7.1"
   s.date          = %q{2019-04-18}
   s.summary       = %q{Set of tools to make Sorbet work with Rails seamlessly.}
   s.authors       = ["Chan Zuckerberg Initiative"]

data/spec/generators/rails-template.rb

@@ -91,6 +91,8 @@ def create_models
         biology: 1,
         dark_art: 999,
       }
+
+      scope :recent, -> { where('created_at > ?', 1.month.ago) }
     end
   RUBY
 
@@ -133,7 +135,7 @@ def create_models
     class Wizard < ApplicationRecord
       validates :name, length: { minimum: 5 }, presence: true
       # simulate conditional validation
-      validates :parent_email, presence: true, if: -> { false }
+      validates :parent_email, presence: true, if: :Slytherin?
 
       typed_enum house: {
         Gryffindor: 0,
@@ -478,11 +480,6 @@ def create_jobs
 end
 
 def add_sorbet_test_files
-  file "typed-override.yaml", <<~YAML
-    true:
-    - ./sorbet_test_cases.rb
-  YAML
-
   copy_file "./sorbet_test_cases.rb", "sorbet_test_cases.rb"
 end
 

data/spec/generators/sorbet_test_cases.rb

@@ -29,26 +29,6 @@ T.assert_type!(T.must(wizard.wand).wizard, Wizard)
 T.assert_type!(wizard.spell_books, SpellBook::ActiveRecord_Associations_CollectionProxy)
 T.assert_type!(wizard.spell_book_ids, T::Array[Integer])
 
-# -- model relation
-# default
-T.assert_type!(Wizard.all, Wizard::ActiveRecord_Relation)
-
-# custom scope
-T.assert_type!(Wizard.recent, Wizard::ActiveRecord_Relation)
-
-# enum scope
-T.assert_type!(Wizard.Gryffindor, Wizard::ActiveRecord_Relation)
-
-# ActiveRecord Querying
-T.assert_type!(Wizard.Gryffindor.recent, Wizard::ActiveRecord_Relation)
-T.assert_type!(Wizard.Gryffindor.recent.unscoped, Wizard::ActiveRecord_Relation)
-T.assert_type!(Wizard.where(id: 1), Wizard::ActiveRecord_Relation)
-T.assert_type!(Wizard.where(id: 1).recent, Wizard::ActiveRecord_Relation)
-T.assert_type!(Wizard.where.not(id: 1), Wizard::ActiveRecord_Relation)
-T.assert_type!(Wizard.preload(:spell_books), Wizard::ActiveRecord_Relation)
-T.assert_type!(Wizard.eager_load(:spell_books), Wizard::ActiveRecord_Relation)
-T.assert_type!(Wizard.order(:id), Wizard::ActiveRecord_Relation)
-
 # Finder methods -- Model
 T.assert_type!(Wizard.exists?(name: 'Test'), T::Boolean)
 T.assert_type!(Wizard.find(wizard.id), Wizard)
@@ -83,6 +63,8 @@ Wizard.find_each { |w| T.assert_type!(w, Wizard) }
 T.assert_type!(Wizard.find_each, T::Enumerator[Wizard])
 Wizard.find_in_batches { |w| T.assert_type!(w, T::Array[Wizard]) }
 T.assert_type!(Wizard.find_in_batches, T::Enumerator[T::Array[Wizard]])
+Wizard.in_batches { |w| T.assert_type!(w, Wizard::ActiveRecord_Relation) }
+T.assert_type!(Wizard.in_batches, ActiveRecord::Batches::BatchEnumerator)
 # T.assert_type!(Wizard.destroy_all, T::Array[Wizard]) # Ignored until we add support
 T.assert_type!(Wizard.any?, T::Boolean)
 T.assert_type!(Wizard.many?, T::Boolean)
@@ -90,6 +72,18 @@ T.assert_type!(Wizard.none?, T::Boolean)
 T.assert_type!(Wizard.one?, T::Boolean)
 # T.assert_type!(Wizard.update_all(name: 'Harry Potter'), Integer) # Ignored until we add support
 # T.assert_type!(Wizard.delete_all, Integer) # Ignored until we add support
+# Query methods
+T.assert_type!(Wizard.all, Wizard::ActiveRecord_Relation)
+T.assert_type!(Wizard.recent, Wizard::ActiveRecord_Relation) # Named scope
+T.assert_type!(Wizard.Gryffindor, Wizard::ActiveRecord_Relation) # Enum scope
+T.assert_type!(Wizard.Gryffindor.recent, Wizard::ActiveRecord_Relation)
+T.assert_type!(Wizard.Gryffindor.recent.unscoped, Wizard::ActiveRecord_Relation)
+T.assert_type!(Wizard.where(id: 1), Wizard::ActiveRecord_Relation)
+T.assert_type!(Wizard.where(id: 1).recent, Wizard::ActiveRecord_Relation)
+T.assert_type!(Wizard.where.not(id: 1), Wizard::ActiveRecord_Relation)
+T.assert_type!(Wizard.preload(:spell_books), Wizard::ActiveRecord_Relation)
+T.assert_type!(Wizard.eager_load(:spell_books), Wizard::ActiveRecord_Relation)
+T.assert_type!(Wizard.order(:id), Wizard::ActiveRecord_Relation)
 
 # Finder methods -- ActiveRecord::Relation
 T.assert_type!(Wizard.all.exists?(name: 'Harry Potter'), T::Boolean)
@@ -126,6 +120,8 @@ Wizard.all.find_each { |w| T.assert_type!(w, Wizard) }
 T.assert_type!(Wizard.all.find_each, T::Enumerator[Wizard])
 Wizard.all.find_in_batches { |w| T.assert_type!(w, T::Array[Wizard]) }
 T.assert_type!(Wizard.all.find_in_batches, T::Enumerator[T::Array[Wizard]])
+Wizard.all.in_batches { |w| T.assert_type!(w, Wizard::ActiveRecord_Relation) }
+T.assert_type!(Wizard.all.in_batches, T::Enumerable[Wizard::ActiveRecord_Relation])
 # T.assert_type!(Wizard.all.destroy_all, T::Array[Wizard]) # Ignored until we add support
 T.assert_type!(Wizard.all.any?, T::Boolean)
 T.assert_type!(Wizard.all.many?, T::Boolean)
@@ -133,6 +129,18 @@ T.assert_type!(Wizard.all.none?, T::Boolean)
 T.assert_type!(Wizard.all.one?, T::Boolean)
 # T.assert_type!(Wizard.all.update_all(name: 'Harry Potter'), Integer) # Ignored until we add support
 # T.assert_type!(Wizard.all.delete_all, Integer) # Ignored until we add support
+# Query methods
+T.assert_type!(Wizard.all.all, Wizard::ActiveRecord_Relation)
+T.assert_type!(Wizard.all.recent, Wizard::ActiveRecord_Relation) # Named scope
+T.assert_type!(Wizard.all.Gryffindor, Wizard::ActiveRecord_Relation) # Enum scope
+T.assert_type!(Wizard.all.Gryffindor.recent, Wizard::ActiveRecord_Relation)
+T.assert_type!(Wizard.all.Gryffindor.recent.unscoped, Wizard::ActiveRecord_Relation)
+T.assert_type!(Wizard.all.where(id: 1), Wizard::ActiveRecord_Relation)
+T.assert_type!(Wizard.all.where(id: 1).recent, Wizard::ActiveRecord_Relation)
+T.assert_type!(Wizard.all.where.not(id: 1), Wizard::ActiveRecord_Relation)
+T.assert_type!(Wizard.all.preload(:spell_books), Wizard::ActiveRecord_Relation)
+T.assert_type!(Wizard.all.eager_load(:spell_books), Wizard::ActiveRecord_Relation)
+T.assert_type!(Wizard.all.order(:id), Wizard::ActiveRecord_Relation)
 # Enumerable methods
 Wizard.all.each { |w| T.assert_type!(w, Wizard) }
 Wizard.all.map { |w| T.assert_type!(w, Wizard) }
@@ -172,10 +180,12 @@ T.assert_type!(spell_books.create!(name: 'Fantastic Beasts') { |s| T.assert_type
 T.assert_type!(spell_books.first_or_create(name: 'Fantastic Beasts') { |s| T.assert_type!(s, SpellBook) }, SpellBook)
 T.assert_type!(spell_books.first_or_create!(name: 'Fantastic Beasts') { |s| T.assert_type!(s, SpellBook) }, SpellBook)
 T.assert_type!(spell_books.first_or_initialize { |s| T.assert_type!(s, SpellBook) }, SpellBook)
-# spell_books.find_each { |s| T.assert_type!(s, SpellBook) } # TODO: Handle for Rails 6
-# T.assert_type!(spell_books.find_each, T::Enumerator[SpellBook]) # TODO: Handle for Rails 6
-# spell_books.find_in_batches { |s| T.assert_type!(s, T::Array[SpellBook]) } # TODO: Handle for Rails 6
-# T.assert_type!(spell_books.find_in_batches, T::Enumerator[T::Array[SpellBook]]) # TODO: Handle for Rails 6
+spell_books.find_each { |s| T.assert_type!(s, SpellBook) }
+T.assert_type!(spell_books.find_each, T::Enumerator[SpellBook])
+spell_books.find_in_batches { |s| T.assert_type!(s, T::Array[SpellBook]) }
+T.assert_type!(spell_books.find_in_batches, T::Enumerator[T::Array[SpellBook]])
+spell_books.in_batches { |s| T.assert_type!(s, SpellBook::ActiveRecord_AssociationRelation) }
+T.assert_type!(spell_books.in_batches, T::Enumerable[SpellBook::ActiveRecord_AssociationRelation])
 # T.assert_type!(spell_books.destroy_all, T::Array[SpellBook]) # Ignored until we add support
 T.assert_type!(spell_books.any?, T::Boolean)
 T.assert_type!(spell_books.many?, T::Boolean)
@@ -183,13 +193,18 @@ T.assert_type!(spell_books.none?, T::Boolean)
 T.assert_type!(spell_books.one?, T::Boolean)
 # T.assert_type!(spell_books.update_all(name: 'Fantastic Beasts'), Integer) # Ignored until we add support
 # T.assert_type!(spell_books.delete_all, Integer) # Ignored until we add support
-# CollectionProxy query also typed correctly!
+# Query methods
+T.assert_type!(spell_books.all, SpellBook::ActiveRecord_AssociationRelation)
+T.assert_type!(spell_books.recent, SpellBook::ActiveRecord_AssociationRelation) # Named scope
+T.assert_type!(spell_books.unclassified, SpellBook::ActiveRecord_AssociationRelation) # Enum scope
+T.assert_type!(spell_books.unclassified.recent, SpellBook::ActiveRecord_AssociationRelation)
+T.assert_type!(spell_books.unclassified.recent.unscoped, SpellBook::ActiveRecord_Relation) # Turns back into relation
 T.assert_type!(spell_books.where(id: 1), SpellBook::ActiveRecord_AssociationRelation)
+T.assert_type!(spell_books.where(id: 1).recent, SpellBook::ActiveRecord_AssociationRelation)
+T.assert_type!(spell_books.where.not(id: 1), SpellBook::ActiveRecord_AssociationRelation)
 T.assert_type!(spell_books.preload(:wizard), SpellBook::ActiveRecord_AssociationRelation)
 T.assert_type!(spell_books.eager_load(:wizard), SpellBook::ActiveRecord_AssociationRelation)
 T.assert_type!(spell_books.order(:id), SpellBook::ActiveRecord_AssociationRelation)
-T.assert_type!(spell_books.where.not(id: 1), SpellBook::ActiveRecord_AssociationRelation)
-T.assert_type!(spell_books.biology, SpellBook::ActiveRecord_AssociationRelation)
 # Enumerable methods
 spell_books.each { |s| T.assert_type!(s, SpellBook) }
 spell_books.map { |s| T.assert_type!(s, SpellBook) }
@@ -237,10 +252,12 @@ T.assert_type!(spell_books_query.create!(name: 'Fantastic Beasts') { |s| T.asser
 T.assert_type!(spell_books_query.first_or_create(name: 'Fantastic Beasts') { |s| T.assert_type!(s, SpellBook) }, SpellBook) # Ignored until we add support
 T.assert_type!(spell_books_query.first_or_create!(name: 'Fantastic Beasts') { |s| T.assert_type!(s, SpellBook) }, SpellBook) # Ignored until we add support
 T.assert_type!(spell_books_query.first_or_initialize { |s| T.assert_type!(s, SpellBook) }, SpellBook) # Ignored until we add support
-# spell_books_query.find_each { |s| T.assert_type!(s, SpellBook) } # TODO: Handle for Rails 6
-# T.assert_type!(spell_books_query.find_each, T::Enumerator[SpellBook]) # TODO: Handle for Rails 6
-# spell_books_query.find_in_batches { |s| T.assert_type!(s, T::Array[SpellBook]) } # TODO: Handle for Rails 6
-# T.assert_type!(spell_books_query.find_in_batches, T::Enumerator[T::Array[SpellBook]]) # TODO: Handle for Rails 6
+spell_books_query.find_each { |s| T.assert_type!(s, SpellBook) }
+T.assert_type!(spell_books_query.find_each, T::Enumerator[SpellBook])
+spell_books_query.find_in_batches { |s| T.assert_type!(s, T::Array[SpellBook]) }
+T.assert_type!(spell_books_query.find_in_batches, T::Enumerator[T::Array[SpellBook]])
+spell_books_query.in_batches { |s| T.assert_type!(s, SpellBook::ActiveRecord_AssociationRelation) }
+T.assert_type!(spell_books_query.in_batches, T::Enumerable[SpellBook::ActiveRecord_AssociationRelation])
 # T.assert_type!(spell_books_query.destroy_all, T::Array[SpellBook]) # Ignored until we add support
 T.assert_type!(spell_books_query.any?, T::Boolean)
 T.assert_type!(spell_books_query.many?, T::Boolean)
@@ -248,12 +265,18 @@ T.assert_type!(spell_books_query.none?, T::Boolean)
 T.assert_type!(spell_books_query.one?, T::Boolean)
 # T.assert_type!(spell_books_query.update_all(name: 'Fantastic Beasts'), Integer) # Ignored until we add support
 # T.assert_type!(spell_books_query.delete_all, Integer) # Ignored until we add support
-# Query chaining
+# Query methods
+T.assert_type!(spell_books_query.all, SpellBook::ActiveRecord_AssociationRelation)
+T.assert_type!(spell_books_query.recent, SpellBook::ActiveRecord_AssociationRelation) # Named scope
+T.assert_type!(spell_books_query.unclassified, SpellBook::ActiveRecord_AssociationRelation) # Enum scope
+T.assert_type!(spell_books_query.unclassified.recent, SpellBook::ActiveRecord_AssociationRelation)
+T.assert_type!(spell_books_query.unclassified.recent.unscoped, SpellBook::ActiveRecord_Relation) # Turns back into relation
+T.assert_type!(spell_books_query.where(id: 1), SpellBook::ActiveRecord_AssociationRelation)
+T.assert_type!(spell_books_query.where(id: 1).recent, SpellBook::ActiveRecord_AssociationRelation)
+T.assert_type!(spell_books_query.where.not(id: 1), SpellBook::ActiveRecord_AssociationRelation)
 T.assert_type!(spell_books_query.preload(:wizard), SpellBook::ActiveRecord_AssociationRelation)
 T.assert_type!(spell_books_query.eager_load(:wizard), SpellBook::ActiveRecord_AssociationRelation)
 T.assert_type!(spell_books_query.order(:id), SpellBook::ActiveRecord_AssociationRelation)
-T.assert_type!(spell_books_query.where.not(id: 1), SpellBook::ActiveRecord_AssociationRelation)
-T.assert_type!(spell_books_query.biology, SpellBook::ActiveRecord_AssociationRelation)
 # Enumerable methods
 spell_books_query.each { |s| T.assert_type!(s, SpellBook) }
 spell_books_query.map { |s| T.assert_type!(s, SpellBook) }
@@ -338,44 +361,6 @@ params = ActionController::Parameters.new({
 typed_params = TypedParams[MyActionParams].new.extract!(params)
 T.assert_type!(typed_params, MyActionParams)
 
-# -- require_typed
-T.assert_type!(
-  params.require_typed(:age, TA[Integer].new),
-  Integer,
-)
-T.assert_type!(
-  params.require_typed(:name, TA[String].new),
-  String,
-)
-info = params.require_typed(:info, TA[ActionController::Parameters].new)
-T.assert_type!(info, ActionController::Parameters)
-T.assert_type!(
-  info.require_typed(:friends, TA[T::Array[String]].new),
-  T::Array[String],
-)
-# -- fetch_typed
-T.assert_type!(
-  params.fetch_typed(:age, TA[Integer].new),
-  Integer,
-)
-T.assert_type!(
-  params.fetch_typed(:name, TA[String].new),
-  String,
-)
-T.assert_type!(
-  params.fetch_typed(:nonexistence, TA[String].new, ''),
-  String,
-)
-T.assert_type!(
-  params.fetch_typed(:nonexistence, TA[T.nilable(String)].new, nil),
-  T.nilable(String),
-)
-T.assert_type!(
-  params.fetch_typed(:nonexistence, TA[T::Array[Integer]].new, []),
-  T::Array[Integer],
-)
-
-
 # -- pluck to tstruct
 class WizardStruct < T::Struct
   const :name, String
@@ -392,23 +377,16 @@ Wizard.all.pluck_to_tstruct(TA[WizardStruct].new).each do |row|
   T.assert_type!(row, WizardStruct)
 end
 
-
 # -- GeneratedUrlHelpers
 class TestHelper
   include GeneratedUrlHelpers
-
-  # need to implement this for the url
-  def default_url_options
-    {
-      protocol: 'http',
-      host: 'localhost',
-      port: 3000,
-    }
-  end
-
-  def test_url_helper
-    T.assert_type!(test_index_path, String)
-    T.assert_type!(test_index_url, String)
-  end
 end
-TestHelper.new.test_url_helper
+
+T.assert_type!(TestHelper.new.test_index_path, String)
+# need to set this config for _url methods
+Rails.application.routes.default_url_options = {
+  protocol: 'http',
+  host: 'localhost',
+  port: 3000,
+}
+T.assert_type!(TestHelper.new.test_index_url, String)

data/spec/job_rbi_formatter_spec.rb

@@ -1,5 +1,5 @@
 require 'rails_helper'
-require 'sorbet-rails/mailer_rbi_formatter'
+require 'sorbet-rails/job_rbi_formatter'
 
 RSpec.describe SorbetRails::MailerRbiFormatter do
 

data/spec/pluck_to_tstruct_spec.rb

@@ -9,6 +9,18 @@ RSpec.describe SorbetRails::PluckToTStruct do
     )
   end
 
+  let!(:harrys_wand) do
+    Wand.create!(
+      wizard: harry,
+      core_type: :phoenix_feather,
+      wood_type: "Holly",
+      chosen_at_date: Date.parse('2019-09-01'),
+      chosen_at_time: Time.parse('2019-09-01T09:00:00Z'),
+      broken: true,
+      broken_at: Time.parse('2019-09-05T15:30:00Z'),
+    )
+  end
+
   let!(:hermione) do
     Wizard.create!(
       name: 'Hermione Granger',
@@ -16,6 +28,18 @@ RSpec.describe SorbetRails::PluckToTStruct do
     )
   end
 
+  let!(:hermiones_wand) do
+    Wand.create!(
+      wizard: hermione,
+      core_type: :phoenix_feather,
+      wood_type: "Vine",
+      chosen_at_date: Date.parse('2019-09-01'),
+      chosen_at_time: Time.parse('2019-09-01T09:00:00Z'),
+      broken: true,
+      broken_at: Time.parse('2019-09-05T15:30:00Z'),
+    )
+  end
+
   class WizardName < T::Struct
     const :name, String
 
@@ -43,6 +67,21 @@ RSpec.describe SorbetRails::PluckToTStruct do
     end
   end
 
+  class WizardWithWandT < T::Struct
+    const :name, String
+    const :house, String
+    const :wand_wood_type, String
+
+    def ==(other)
+      return false unless other.is_a?(self.class)
+      name == other.name && house == other.house
+    end
+
+    def eql?(other)
+      self == other
+    end
+  end
+
   shared_examples 'pluck_to_tstruct' do |struct_type, expected_values|
     it 'plucks correctly from ActiveRecord model' do
       plucked = Wizard.pluck_to_tstruct(TA[struct_type].new)
@@ -55,6 +94,18 @@ RSpec.describe SorbetRails::PluckToTStruct do
     end
   end
 
+  shared_examples 'pluck_to_tstruct with associations' do |struct_type, associations, expected_values|
+    it 'plucks correctly from ActiveRecord model with associations' do
+      plucked = Wizard.joins(:wand).pluck_to_tstruct(TA[struct_type].new, associations: associations)
+      expect(plucked).to match_array(expected_values)
+    end
+
+    it 'plucks correctly from ActiveRecord relation with associations' do
+      plucked = Wizard.all.joins(:wand).pluck_to_tstruct(TA[struct_type].new, associations: associations)
+      expect(plucked).to match_array(expected_values)
+    end
+  end
+
   context 'pluck 1 attribute' do
     it_should_behave_like 'pluck_to_tstruct', WizardName, [
       WizardName.new(name: "Harry Potter"),
@@ -72,8 +123,30 @@ RSpec.describe SorbetRails::PluckToTStruct do
   context 'given a wrong type' do
     it 'should raise error' do
       expect {
-        plucked = Wizard.pluck_to_tstruct(TA[String].new)
+        Wizard.pluck_to_tstruct(TA[String].new)
       }.to raise_error(SorbetRails::PluckToTStruct::UnexpectedType)
     end
   end
+
+  context "given associations mappings that don't exist in T::Struct" do
+    it 'should raise error' do
+      expect {
+        Wizard.pluck_to_tstruct(
+          TA[WizardWithWandT].new,
+          associations: { wood_type: "wands.wood_type" }
+        )
+      }.to raise_error(SorbetRails::PluckToTStruct::UnexpectedAssociations)
+    end
+  end
+
+  context 'pluck with associations' do
+    associations = { wand_wood_type: "wands.wood_type" }
+
+    expected = [
+      WizardWithWandT.new(name: "Harry Potter", house: "Gryffindor", wand_wood_type: "Holly"),
+      WizardWithWandT.new(name: "Hermione Granger", house: "Gryffindor", wand_wood_type: "Vine"),
+    ]
+
+    it_should_behave_like 'pluck_to_tstruct with associations', WizardWithWandT, associations, expected
+  end
 end