WIP

Author: byroot

activerecord/lib/active_record/associations/association.rb

@@ -191,7 +191,7 @@ def async_load_target
         @target = find_target(async: true) if (@stale_state && stale_target?) || find_target?
 
         loaded! unless loaded?
-        target
+        @target
       end
 
       # We can't dump @reflection and @through_reflection since it contains the scope proc

activerecord/test/cases/associations/belongs_to_associations_test.rb

@@ -1843,19 +1843,35 @@ def test_destroy_linked_models
 end
 
 class AsyncBelongsToAssociationsTest < ActiveRecord::TestCase
+  include WaitForAsyncTestHelper
+
   fixtures :companies
 
   self.use_transactional_tests = false
 
   def test_temp_async_load_belongs_to
-    # TODO: proper test?
     client = Client.find(3)
     first_firm = companies(:first_firm)
-    assert_queries_match(/LIMIT|ROWNUM <=|FETCH FIRST/) do
-      client.association(:firm).async_load_target
 
+    promise = client.association(:firm).async_load_target
+    wait_for_async_query
+
+    events = []
+    callback = -> (event) do
+      events << event unless event.payload[:name] == "SCHEMA"
+    end
+    ActiveSupport::Notifications.subscribed(callback, "sql.active_record") do
+      client.firm
+    end
+
+    assert_no_queries do
       assert_equal first_firm, client.firm
       assert_equal first_firm.name, client.firm.name
     end
+
+    assert_equal 1, events.size
+    assert_equal true, events.first.payload[:async]
+  ensure
+    ActiveRecord::Base.logger = Logger.new(nil)
   end
 end

activerecord/test/cases/helper.rb

@@ -40,36 +40,53 @@
 ActiveRecord::ConnectionAdapters.register("abstract", "ActiveRecord::ConnectionAdapters::AbstractAdapter", "active_record/connection_adapters/abstract_adapter")
 ActiveRecord::ConnectionAdapters.register("fake", "FakeActiveRecordAdapter", File.expand_path("../support/fake_adapter.rb", __dir__))
 
-class SQLSubscriber
-  attr_reader :logged
-  attr_reader :payloads
+class ActiveRecord::TestCase
+  class SQLSubscriber
+    attr_reader :logged
+    attr_reader :payloads
+
+    def initialize
+      @logged = []
+      @payloads = []
+    end
+
+    def start(name, id, payload)
+      @payloads << payload
+      @logged << [payload[:sql].squish, payload[:name], payload[:binds]]
+    end
 
-  def initialize
-    @logged = []
-    @payloads = []
+    def finish(name, id, payload); end
   end
 
-  def start(name, id, payload)
-    @payloads << payload
-    @logged << [payload[:sql].squish, payload[:name], payload[:binds]]
+  module InTimeZone
+    private
+      def in_time_zone(zone)
+        old_zone  = Time.zone
+        old_tz    = ActiveRecord::Base.time_zone_aware_attributes
+
+        Time.zone = zone ? ActiveSupport::TimeZone[zone] : nil
+        ActiveRecord::Base.time_zone_aware_attributes = !zone.nil?
+        yield
+      ensure
+        Time.zone = old_zone
+        ActiveRecord::Base.time_zone_aware_attributes = old_tz
+      end
   end
 
-  def finish(name, id, payload); end
-end
+  module WaitForAsyncTestHelper
+    private
+      def wait_for_async_query(connection = ActiveRecord::Base.lease_connection, timeout: 5)
+        return unless connection.async_enabled?
 
-module InTimeZone
-  private
-    def in_time_zone(zone)
-      old_zone  = Time.zone
-      old_tz    = ActiveRecord::Base.time_zone_aware_attributes
-
-      Time.zone = zone ? ActiveSupport::TimeZone[zone] : nil
-      ActiveRecord::Base.time_zone_aware_attributes = !zone.nil?
-      yield
-    ensure
-      Time.zone = old_zone
-      ActiveRecord::Base.time_zone_aware_attributes = old_tz
-    end
+        executor = connection.pool.async_executor
+        (timeout * 100).times do
+          return unless executor.scheduled_task_count > executor.completed_task_count
+          sleep 0.01
+        end
+
+        raise Timeout::Error, "The async executor wasn't drained after #{timeout} seconds"
+      end
+  end
 end
 
 # Encryption

activerecord/test/cases/relation/load_async_test.rb

@@ -7,21 +7,6 @@
 require "models/other_dog"
 
 module ActiveRecord
-  module WaitForAsyncTestHelper
-    private
-      def wait_for_async_query(connection = ActiveRecord::Base.lease_connection, timeout: 5)
-        return unless connection.async_enabled?
-
-        executor = connection.pool.async_executor
-        (timeout * 100).times do
-          return unless executor.scheduled_task_count > executor.completed_task_count
-          sleep 0.01
-        end
-
-        raise Timeout::Error, "The async executor wasn't drained after #{timeout} seconds"
-      end
-  end
-
   class LoadAsyncTest < ActiveRecord::TestCase
     include WaitForAsyncTestHelper