From 5c562a6351ba12ed7a33d5959121b78cbd5c132b Mon Sep 17 00:00:00 2001
From: Eugene Kenny <elkenny@gmail.com>
Date: Thu, 29 Oct 2020 21:15:15 +0000
Subject: [PATCH] Only replace existing pools in transactional tests

`setup_shared_connection_pool` replaces reading connection pools with
writing ones, so that replica reads work as they would in production.

If a shard doesn't have a reading connection we should skip it, since
adding one could lead to code that only works in the test environment.
---
 .../lib/active_record/test_fixtures.rb        |  1 +
 activerecord/test/cases/fixtures_test.rb      | 30 +++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/activerecord/lib/active_record/test_fixtures.rb b/activerecord/lib/active_record/test_fixtures.rb
index 7d533ab0ee597..98bc2584a877a 100644
--- a/activerecord/lib/active_record/test_fixtures.rb
+++ b/activerecord/lib/active_record/test_fixtures.rb
@@ -214,6 +214,7 @@ def setup_shared_connection_pool
             pool_manager.shard_names.each do |shard_name|
               writing_pool_config = pool_manager.get_pool_config(ActiveRecord::Base.writing_role, shard_name)
               pool_manager.role_names.each do |role|
+                next unless pool_manager.get_pool_config(role, shard_name)
                 pool_manager.set_pool_config(role, shard_name, writing_pool_config)
               end
             end
diff --git a/activerecord/test/cases/fixtures_test.rb b/activerecord/test/cases/fixtures_test.rb
index 0c0724aa979a0..19ce2b427301c 100644
--- a/activerecord/test/cases/fixtures_test.rb
+++ b/activerecord/test/cases/fixtures_test.rb
@@ -1453,6 +1453,21 @@ def test_writing_and_reading_connections_are_the_same
       assert_equal rw_conn, ro_conn
     end
 
+    def test_only_existing_connections_are_replaced
+      ActiveRecord::Base.connects_to shards: {
+        default: { writing: :default, reading: :readonly },
+        two: { writing: :default }
+      }
+
+      setup_shared_connection_pool
+
+      assert_raises(ActiveRecord::ConnectionNotEstablished) do
+        ActiveRecord::Base.connected_to(role: :reading, shard: :two) do
+          ActiveRecord::Base.retrieve_connection
+        end
+      end
+    end
+
     private
       def config
         { "default" => default_config, "readonly" => readonly_config }
@@ -1516,6 +1531,21 @@ def test_writing_and_reading_connections_are_the_same_with_legacy_handling
       assert_equal rw_conn, ro_conn
     end
 
+    def test_only_existing_connections_are_replaced
+      ActiveRecord::Base.connects_to shards: {
+        default: { writing: :default, reading: :readonly },
+        two: { writing: :default }
+      }
+
+      setup_shared_connection_pool
+
+      assert_raises(ActiveRecord::ConnectionNotEstablished) do
+        ActiveRecord::Base.connected_to(role: :reading, shard: :two) do
+          ActiveRecord::Base.retrieve_connection
+        end
+      end
+    end
+
     private
       def config
         { "default" => default_config, "readonly" => readonly_config }