1 files changed, 43 insertions, 3 deletions
diff --git a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
index c647f5a8f5d..86fae14b961 100644
--- a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
@@ -5,7 +5,12 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
   let(:conflict_error) { Class.new(RuntimeError) }
   let(:db_host) { ActiveRecord::Base.connection_pool.db_config.host }
-  let(:lb) { described_class.new([db_host, db_host]) }
+  let(:config) do
+    Gitlab::Database::LoadBalancing::Configuration
+      .new(ActiveRecord::Base, [db_host, db_host])
+  end
+
+  let(:lb) { described_class.new(config) }
   let(:request_cache) { lb.send(:request_cache) }
 
   before do
@@ -41,6 +46,19 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
     top_error
   end
 
+  describe '#initialize' do
+    it 'ignores the hosts when the primary_only option is enabled' do
+      config = Gitlab::Database::LoadBalancing::Configuration
+        .new(ActiveRecord::Base, [db_host])
+      lb = described_class.new(config, primary_only: true)
+      hosts = lb.host_list.hosts
+
+      expect(hosts.length).to eq(1)
+      expect(hosts.first)
+        .to be_instance_of(Gitlab::Database::LoadBalancing::PrimaryHost)
+    end
+  end
+
   describe '#read' do
     it 'yields a connection for a read' do
       connection = double(:connection)
@@ -121,6 +139,19 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
       expect { |b| lb.read(&b) }
         .to yield_with_args(ActiveRecord::Base.retrieve_connection)
     end
+
+    it 'uses the primary when the primary_only option is enabled' do
+      config = Gitlab::Database::LoadBalancing::Configuration
+        .new(ActiveRecord::Base)
+      lb = described_class.new(config, primary_only: true)
+
+      # When no hosts are configured, we don't want to produce any warnings, as
+      # they aren't useful/too noisy.
+      expect(Gitlab::Database::LoadBalancing::Logger).not_to receive(:warn)
+
+      expect { |b| lb.read(&b) }
+        .to yield_with_args(ActiveRecord::Base.retrieve_connection)
+    end
   end
 
   describe '#read_write' do
@@ -152,8 +183,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
     end
 
     it 'does not create conflicts with other load balancers when caching hosts' do
-      lb1 = described_class.new([db_host, db_host], ActiveRecord::Base)
-      lb2 = described_class.new([db_host, db_host], Ci::CiDatabaseRecord)
+      ci_config = Gitlab::Database::LoadBalancing::Configuration
+        .new(Ci::CiDatabaseRecord, [db_host, db_host])
+
+      lb1 = described_class.new(config)
+      lb2 = described_class.new(ci_config)
 
       host1 = lb1.host
       host2 = lb2.host
@@ -283,6 +317,12 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
 
       expect(lb.connection_error?(error)).to eq(false)
     end
+
+    it 'returns false for ActiveRecord errors without a cause' do
+      error = ActiveRecord::RecordNotUnique.new
+
+      expect(lb.connection_error?(error)).to eq(false)
+    end
   end
 
   describe '#serialization_failure?' do