1 files changed, 90 insertions, 18 deletions
diff --git a/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb b/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
index a27341a3324..e9bc465b1c7 100644
--- a/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
@@ -3,13 +3,18 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
-  let(:load_balancer) { Gitlab::Database::LoadBalancing::LoadBalancer.new([]) }
+  let(:load_balancer) do
+    Gitlab::Database::LoadBalancing::LoadBalancer.new(
+      Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base)
+    )
+  end
+
   let(:service) do
     described_class.new(
+      load_balancer,
       nameserver: 'localhost',
       port: 8600,
-      record: 'foo',
-      load_balancer: load_balancer
+      record: 'foo'
     )
   end
 
@@ -26,11 +31,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
     describe ':record_type' do
       subject do
         described_class.new(
+          load_balancer,
           nameserver: 'localhost',
           port: 8600,
           record: 'foo',
-          record_type: record_type,
-          load_balancer: load_balancer
+          record_type: record_type
         )
       end
 
@@ -69,18 +74,69 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
   end
 
   describe '#perform_service_discovery' do
-    it 'reports exceptions to Sentry' do
-      error = StandardError.new
+    context 'without any failures' do
+      it 'runs once' do
+        expect(service)
+          .to receive(:refresh_if_necessary).once
+
+        expect(service).not_to receive(:sleep)
+
+        expect(Gitlab::ErrorTracking).not_to receive(:track_exception)
+
+        service.perform_service_discovery
+      end
+    end
+    context 'with failures' do
+      before do
+        allow(Gitlab::ErrorTracking).to receive(:track_exception)
+        allow(service).to receive(:sleep)
+      end
+
+      let(:valid_retry_sleep_duration) { satisfy { |val| described_class::RETRY_DELAY_RANGE.include?(val) } }
+
+      it 'retries service discovery when under the retry limit' do
+        error = StandardError.new
+
+        expect(service)
+          .to receive(:refresh_if_necessary)
+          .and_raise(error).exactly(described_class::MAX_DISCOVERY_RETRIES - 1).times.ordered
+
+        expect(service)
+          .to receive(:sleep).with(valid_retry_sleep_duration)
+          .exactly(described_class::MAX_DISCOVERY_RETRIES - 1).times
+
+        expect(service).to receive(:refresh_if_necessary).and_return(45).ordered
+
+        expect(service.perform_service_discovery).to eq(45)
+      end
+
+      it 'does not retry service discovery after exceeding the limit' do
+        error = StandardError.new
+
+        expect(service)
+          .to receive(:refresh_if_necessary)
+          .and_raise(error).exactly(described_class::MAX_DISCOVERY_RETRIES).times
+
+        expect(service)
+          .to receive(:sleep).with(valid_retry_sleep_duration)
+          .exactly(described_class::MAX_DISCOVERY_RETRIES).times
+
+        service.perform_service_discovery
+      end
 
-      expect(service)
-        .to receive(:refresh_if_necessary)
-        .and_raise(error)
+      it 'reports exceptions to Sentry' do
+        error = StandardError.new
+
+        expect(service)
+          .to receive(:refresh_if_necessary)
+                .and_raise(error).exactly(described_class::MAX_DISCOVERY_RETRIES).times
 
-      expect(Gitlab::ErrorTracking)
-        .to receive(:track_exception)
-        .with(error)
+        expect(Gitlab::ErrorTracking)
+          .to receive(:track_exception)
+                .with(error).exactly(described_class::MAX_DISCOVERY_RETRIES).times
 
-      service.perform_service_discovery
+        service.perform_service_discovery
+      end
     end
   end
 
@@ -133,7 +189,10 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
     let(:address_bar) { described_class::Address.new('bar') }
 
     let(:load_balancer) do
-      Gitlab::Database::LoadBalancing::LoadBalancer.new([address_foo])
+      Gitlab::Database::LoadBalancing::LoadBalancer.new(
+        Gitlab::Database::LoadBalancing::Configuration
+          .new(ActiveRecord::Base, [address_foo])
+      )
     end
 
     before do
@@ -166,11 +225,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
   describe '#addresses_from_dns' do
     let(:service) do
       described_class.new(
+        load_balancer,
         nameserver: 'localhost',
         port: 8600,
         record: 'foo',
-        record_type: record_type,
-        load_balancer: load_balancer
+        record_type: record_type
       )
     end
 
@@ -224,6 +283,16 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
         expect(service.addresses_from_dns).to eq([90, addresses])
       end
     end
+
+    context 'when the resolver returns an empty response' do
+      let(:packet) { double(:packet, answer: []) }
+
+      let(:record_type) { 'A' }
+
+      it 'raises EmptyDnsResponse' do
+        expect { service.addresses_from_dns }.to raise_error(Gitlab::Database::LoadBalancing::ServiceDiscovery::EmptyDnsResponse)
+      end
+    end
   end
 
   describe '#new_wait_time_for' do
@@ -246,7 +315,10 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
 
   describe '#addresses_from_load_balancer' do
     let(:load_balancer) do
-      Gitlab::Database::LoadBalancing::LoadBalancer.new(%w[b a])
+      Gitlab::Database::LoadBalancing::LoadBalancer.new(
+        Gitlab::Database::LoadBalancing::Configuration
+          .new(ActiveRecord::Base, %w[b a])
+      )
     end
 
     it 'returns the ordered host names of the load balancer' do