2 files changed, 32 insertions, 3 deletions
diff --git a/spec/lib/gitlab/database/count/reltuples_count_strategy_spec.rb b/spec/lib/gitlab/database/count/reltuples_count_strategy_spec.rb
index b44e8c5a110..bd3c66d0548 100644
--- a/spec/lib/gitlab/database/count/reltuples_count_strategy_spec.rb
+++ b/spec/lib/gitlab/database/count/reltuples_count_strategy_spec.rb
@@ -6,10 +6,11 @@ describe Gitlab::Database::Count::ReltuplesCountStrategy do
     create(:identity)
   end
 
-  let(:models) { [Project, Identity] }
   subject { described_class.new(models).count }
 
   describe '#count', :postgresql do
+    let(:models) { [Project, Identity] }
+
     context 'when reltuples is up to date' do
       before do
         ActiveRecord::Base.connection.execute('ANALYZE projects')
@@ -23,6 +24,22 @@ describe Gitlab::Database::Count::ReltuplesCountStrategy do
       end
     end
 
+    context 'when models using single-type inheritance are used' do
+      let(:models) { [Group, CiService, Namespace] }
+
+      before do
+        models.each do |model|
+          ActiveRecord::Base.connection.execute("ANALYZE #{model.table_name}")
+        end
+      end
+
+      it 'returns nil counts for inherited tables' do
+        models.each { |model| expect(model).not_to receive(:count) }
+
+        expect(subject).to eq({ Namespace => 3 })
+      end
+    end
+
     context 'insufficient permissions' do
       it 'returns an empty hash' do
         allow(ActiveRecord::Base).to receive(:transaction).and_raise(PG::InsufficientPrivilege)
diff --git a/spec/lib/gitlab/database/count/tablesample_count_strategy_spec.rb b/spec/lib/gitlab/database/count/tablesample_count_strategy_spec.rb
index 203f9344a41..40d810b195b 100644
--- a/spec/lib/gitlab/database/count/tablesample_count_strategy_spec.rb
+++ b/spec/lib/gitlab/database/count/tablesample_count_strategy_spec.rb
@@ -4,15 +4,23 @@ describe Gitlab::Database::Count::TablesampleCountStrategy do
   before do
     create_list(:project, 3)
     create(:identity)
+    create(:group)
   end
 
-  let(:models) { [Project, Identity] }
+  let(:models) { [Project, Identity, Group, Namespace] }
   let(:strategy) { described_class.new(models) }
 
   subject { strategy.count }
 
   describe '#count', :postgresql do
-    let(:estimates) { { Project => threshold + 1, Identity => threshold - 1 } }
+    let(:estimates) do
+      {
+        Project => threshold + 1,
+        Identity => threshold - 1,
+        Group => threshold + 1,
+        Namespace => threshold + 1
+      }
+    end
     let(:threshold) { Gitlab::Database::Count::TablesampleCountStrategy::EXACT_COUNT_THRESHOLD }
 
     before do
@@ -30,9 +38,13 @@ describe Gitlab::Database::Count::TablesampleCountStrategy do
     context 'for tables with an estimated large size' do
       it 'performs a tablesample count' do
         expect(Project).not_to receive(:count)
+        expect(Group).not_to receive(:count)
+        expect(Namespace).not_to receive(:count)
 
         result = subject
         expect(result[Project]).to eq(3)
+        expect(result[Group]).to eq(1)
+        expect(result[Namespace]).to eq(4)
       end
     end