1 files changed, 45 insertions, 36 deletions
diff --git a/spec/lib/gitlab/database/count_spec.rb b/spec/lib/gitlab/database/count_spec.rb
index 9d9caaabe16..407d9470785 100644
--- a/spec/lib/gitlab/database/count_spec.rb
+++ b/spec/lib/gitlab/database/count_spec.rb
@@ -3,59 +3,68 @@ require 'spec_helper'
 describe Gitlab::Database::Count do
   before do
     create_list(:project, 3)
+    create(:identity)
   end
 
-  describe '.execute_estimate_if_updated_recently', :postgresql do
-    context 'when reltuples have not been updated' do
-      before do
-        expect(described_class).to receive(:reltuples_updated_recently?).and_return(false)
-      end
+  let(:models) { [Project, Identity] }
 
-      it 'returns nil' do
-        expect(described_class.execute_estimate_if_updated_recently(Project)).to be nil
-      end
-    end
+  describe '.approximate_counts' do
+    context 'with MySQL' do
+      context 'when reltuples have not been updated' do
+        it 'counts all models the normal way' do
+          expect(Gitlab::Database).to receive(:postgresql?).and_return(false)
 
-    context 'when reltuples have been updated' do
-      before do
-        ActiveRecord::Base.connection.execute('ANALYZE projects')
-      end
+          expect(Project).to receive(:count).and_call_original
+          expect(Identity).to receive(:count).and_call_original
 
-      it 'calls postgresql_estimate_query' do
-        expect(described_class).to receive(:postgresql_estimate_query).with(Project).and_call_original
-        expect(described_class.execute_estimate_if_updated_recently(Project)).to eq(3)
+          expect(described_class.approximate_counts(models)).to eq({ Project => 3, Identity => 1 })
+        end
       end
     end
-  end
 
-  describe '.approximate_count' do
-    context 'when reltuples have not been updated' do
-      it 'counts all projects the normal way' do
-        allow(described_class).to receive(:reltuples_updated_recently?).and_return(false)
+    context 'with PostgreSQL', :postgresql do
+      describe 'when reltuples have not been updated' do
+        it 'counts all models the normal way' do
+          expect(described_class).to receive(:reltuples_from_recently_updated).with(%w(projects identities)).and_return({})
 
-        expect(Project).to receive(:count).and_call_original
-        expect(described_class.approximate_count(Project)).to eq(3)
+          expect(Project).to receive(:count).and_call_original
+          expect(Identity).to receive(:count).and_call_original
+          expect(described_class.approximate_counts(models)).to eq({ Project => 3, Identity => 1 })
+        end
       end
-    end
 
-    context 'no permission' do
-      it 'falls back to standard query' do
-        allow(described_class).to receive(:reltuples_updated_recently?).and_raise(PG::InsufficientPrivilege)
+      describe 'no permission' do
+        it 'falls back to standard query' do
+          allow(described_class).to receive(:postgresql_estimate_query).and_raise(PG::InsufficientPrivilege)
 
-        expect(Project).to receive(:count).and_call_original
-        expect(described_class.approximate_count(Project)).to eq(3)
+          expect(Project).to receive(:count).and_call_original
+          expect(Identity).to receive(:count).and_call_original
+          expect(described_class.approximate_counts(models)).to eq({ Project => 3, Identity => 1 })
+        end
       end
-    end
 
-    describe 'when reltuples have been updated', :postgresql do
-      before do
-        ActiveRecord::Base.connection.execute('ANALYZE projects')
+      describe 'when some reltuples have been updated' do
+        it 'counts projects in the fast way' do
+          expect(described_class).to receive(:reltuples_from_recently_updated).with(%w(projects identities)).and_return({ 'projects' => 3 })
+
+          expect(Project).not_to receive(:count).and_call_original
+          expect(Identity).to receive(:count).and_call_original
+          expect(described_class.approximate_counts(models)).to eq({ Project => 3, Identity => 1 })
+        end
       end
 
-      it 'counts all projects in the fast way' do
-        expect(described_class).to receive(:postgresql_estimate_query).with(Project).and_call_original
+      describe 'when all reltuples have been updated' do
+        before do
+          ActiveRecord::Base.connection.execute('ANALYZE projects')
+          ActiveRecord::Base.connection.execute('ANALYZE identities')
+        end
+
+        it 'counts models with the standard way' do
+          expect(Project).not_to receive(:count)
+          expect(Identity).not_to receive(:count)
 
-        expect(described_class.approximate_count(Project)).to eq(3)
+          expect(described_class.approximate_counts(models)).to eq({ Project => 3, Identity => 1 })
+        end
       end
     end
   end