From 3574a1be6c4f02e28beae201160ff13537dc4172 Mon Sep 17 00:00:00 2001
From: Eugene Kenny <elkenny@gmail.com>
Date: Wed, 21 Apr 2021 22:03:19 +0100
Subject: [PATCH] Merge pull request #40384 from
 eugeneius/teardown_shared_connection_pool

Restore connection pools after transactional tests
---
 .../lib/active_record/test_fixtures.rb        | 18 +++++++++++++++++
 activerecord/test/cases/fixtures_test.rb      | 20 +++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/activerecord/lib/active_record/test_fixtures.rb b/activerecord/lib/active_record/test_fixtures.rb
index c374368662512..05c00ba92530e 100644
--- a/activerecord/lib/active_record/test_fixtures.rb
+++ b/activerecord/lib/active_record/test_fixtures.rb
@@ -112,6 +112,8 @@ def setup_fixtures(config = ActiveRecord::Base)
 
       # Load fixtures once and begin transaction.
       if run_in_transaction?
+        @saved_pool_configs = Hash.new { |hash, key| hash[key] = {} }
+
         if @@already_loaded_fixtures[self.class]
           @loaded_fixtures = @@already_loaded_fixtures[self.class]
         else
@@ -166,6 +168,7 @@ def teardown_fixtures
           connection.pool.lock_thread = false
         end
         @fixture_connections.clear
+        teardown_shared_connection_pool
       else
         ActiveRecord::FixtureSet.reset_cache
       end
@@ -195,12 +198,27 @@ def setup_shared_connection_pool
               name = pool.spec.name
               writing_connection = writing_handler.retrieve_connection_pool(name)
               return unless writing_connection
+
+              reading_connection = handler.send(:owner_to_pool)[name]
+              next if reading_connection == writing_connection
+
+              @saved_pool_configs[handler][name] = reading_connection
               handler.send(:owner_to_pool)[name] = writing_connection
             end
           end
         end
       end
 
+      def teardown_shared_connection_pool
+        @saved_pool_configs.each_pair do |handler, pools|
+          pools.each_pair do |name, pool|
+            handler.send(:owner_to_pool)[name] = pool
+          end
+        end
+
+        @saved_pool_configs.clear
+      end
+
       def load_fixtures(config)
         fixtures = ActiveRecord::FixtureSet.create_fixtures(fixture_path, fixture_table_names, fixture_class_names, config)
         Hash[fixtures.map { |f| [f.name, f] }]
diff --git a/activerecord/test/cases/fixtures_test.rb b/activerecord/test/cases/fixtures_test.rb
index bf3a1e76af728..6a1cd61d0fb7e 100644
--- a/activerecord/test/cases/fixtures_test.rb
+++ b/activerecord/test/cases/fixtures_test.rb
@@ -1369,6 +1369,8 @@ def setup
       ActiveRecord::Base.connection_handlers = {}
       ActiveRecord::Base.connection_handler = handler
       ActiveRecord::Base.connects_to(database: { writing: :default, reading: :readonly })
+
+      setup_shared_connection_pool
     end
 
     def teardown
@@ -1392,6 +1394,24 @@ def test_writing_and_reading_connections_are_the_same
       ro_conn = ActiveRecord::Base.connection_handlers[:reading].connection_pool_list.first.connection
 
       assert_equal rw_conn, ro_conn
+
+      teardown_shared_connection_pool
+
+      rw_conn = ActiveRecord::Base.connection_handlers[:writing].connection_pool_list.first.connection
+      ro_conn = ActiveRecord::Base.connection_handlers[:reading].connection_pool_list.first.connection
+
+      assert_not_equal rw_conn, ro_conn
+    end
+
+    def test_only_existing_connections_are_restored
+      ActiveRecord::Base.connection_handlers = { writing: ActiveRecord::Base.default_connection_handler }
+      teardown_shared_connection_pool
+
+      assert_raises(ActiveRecord::ConnectionNotEstablished) do
+        ActiveRecord::Base.connected_to(role: :reading) do
+          ActiveRecord::Base.retrieve_connection
+        end
+      end
     end
 
     private