2 files changed, 60 insertions, 6 deletions
diff --git a/spec/lib/gitlab/database/load_balancing/resolver_spec.rb b/spec/lib/gitlab/database/load_balancing/resolver_spec.rb
index 0051cf50255..4af36693383 100644
--- a/spec/lib/gitlab/database/load_balancing/resolver_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/resolver_spec.rb
@@ -2,15 +2,16 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::Database::LoadBalancing::Resolver do
+RSpec.describe Gitlab::Database::LoadBalancing::Resolver, :freeze_time, feature_category: :database do
   describe '#resolve' do
     let(:ip_addr) { IPAddr.new('127.0.0.2') }
 
     context 'when nameserver is an IP' do
       it 'returns an IPAddr object' do
         service = described_class.new('127.0.0.2')
+        response = service.resolve
 
-        expect(service.resolve).to eq(ip_addr)
+        expect(response.address).to eq(ip_addr)
       end
     end
 
@@ -22,12 +23,14 @@ RSpec.describe Gitlab::Database::LoadBalancing::Resolver do
           allow(instance).to receive(:getaddress).with('localhost').and_return('127.0.0.2')
         end
 
-        expect(subject).to eq(ip_addr)
+        expect(subject.address).to eq(ip_addr)
       end
 
       context 'when nameserver is not in the hosts file' do
+        let(:raw_ttl) { 10 }
+
         it 'looks the nameserver up in DNS' do
-          resource = double(:resource, address: ip_addr)
+          resource = double(:resource, address: ip_addr, ttl: raw_ttl)
           packet = double(:packet, answer: [resource])
 
           allow_next_instance_of(Resolv::Hosts) do |instance|
@@ -38,7 +41,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::Resolver do
             .with('localhost', Net::DNS::A)
             .and_return(packet)
 
-          expect(subject).to eq(ip_addr)
+          expect(subject.address).to eq(ip_addr)
+          expect(subject.ttl).to eq(raw_ttl.seconds.from_now)
         end
 
         context 'when nameserver is not in DNS' do
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 984d60e9962..bfd9c644ffa 100644
--- a/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
@@ -2,7 +2,7 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
+RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery, feature_category: :database do
   let(:load_balancer) do
     configuration = Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base)
     configuration.service_discovery[:record] = 'localhost'
@@ -23,6 +23,8 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
     resource = double(:resource, address: IPAddr.new('127.0.0.1'))
     packet = double(:packet, answer: [resource])
 
+    service.instance_variable_set(:@nameserver_ttl, Gitlab::Database::LoadBalancing::Resolver::FAR_FUTURE_TTL)
+
     allow(Net::DNS::Resolver).to receive(:start)
       .with('localhost', Net::DNS::A)
       .and_return(packet)
@@ -362,4 +364,52 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
       expect(service.addresses_from_load_balancer).to eq(addresses)
     end
   end
+
+  describe '#resolver', :freeze_time do
+    context 'without predefined resolver' do
+      it 'fetches a new resolver and assigns it to the instance variable' do
+        expect(service.instance_variable_get(:@resolver)).not_to be_present
+
+        service_resolver = service.resolver
+
+        expect(service.instance_variable_get(:@resolver)).to be_present
+        expect(service_resolver).to be_present
+      end
+    end
+
+    context 'with predefined resolver' do
+      let(:resolver) do
+        Net::DNS::Resolver.new(
+          nameservers: 'localhost',
+          port: 8600
+        )
+      end
+
+      before do
+        service.instance_variable_set(:@resolver, resolver)
+      end
+
+      context "when nameserver's TTL is in the future" do
+        it 'returns the existing resolver' do
+          expect(service.resolver).to eq(resolver)
+        end
+      end
+
+      context "when nameserver's TTL is in the past" do
+        before do
+          service.instance_variable_set(
+            :@nameserver_ttl,
+            1.minute.ago
+          )
+        end
+
+        it 'fetches new resolver' do
+          service_resolver = service.resolver
+
+          expect(service_resolver).to be_present
+          expect(service_resolver).not_to eq(resolver)
+        end
+      end
+    end
+  end
 end