Fix off-by-one expectation and simplify partitioning specs a little

spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb

@@ -165,10 +165,13 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
       it 'creates a partition spanning over each month in the range given' do
         migration.partition_table_by_date template_table, partition_column, min_date: min_date, max_date: max_date
 
-        expect_range_partition_of("#{partitioned_table}_000000", partitioned_table, 'MINVALUE', "'2019-12-01 00:00:00'")
-        expect_range_partition_of("#{partitioned_table}_201912", partitioned_table, "'2019-12-01 00:00:00'", "'2020-01-01 00:00:00'")
-        expect_range_partition_of("#{partitioned_table}_202001", partitioned_table, "'2020-01-01 00:00:00'", "'2020-02-01 00:00:00'")
-        expect_range_partition_of("#{partitioned_table}_202002", partitioned_table, "'2020-02-01 00:00:00'", "'2020-03-01 00:00:00'")
+        expect_range_partitions_for(partitioned_table, {
+          '000000' => ['MINVALUE', "'2019-12-01 00:00:00'"],
+          '201912' => ["'2019-12-01 00:00:00'", "'2020-01-01 00:00:00'"],
+          '202001' => ["'2020-01-01 00:00:00'", "'2020-02-01 00:00:00'"],
+          '202002' => ["'2020-02-01 00:00:00'", "'2020-03-01 00:00:00'"],
+          '202003' => ["'2020-03-01 00:00:00'", "'2020-04-01 00:00:00'"]
+        })
       end
     end
 

spec/support/helpers/partitioning_helpers.rb

@@ -16,6 +16,21 @@ module PartitioningHelpers
     expect(definition['condition']).to eq("FOR VALUES FROM (#{min_value}) TO (#{max_value})")
   end
 
+  def expect_total_partitions(table_name, count, schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)
+    partitions = find_partitions(table_name, schema: schema)
+
+    expect(partitions.size).to eq(count)
+  end
+
+  def expect_range_partitions_for(table_name, partitions)
+    partitions.each do |suffix, (min_value, max_value)|
+      partition_name = "#{table_name}_#{suffix}"
+      expect_range_partition_of(partition_name, table_name, min_value, max_value)
+    end
+
+    expect_total_partitions(table_name, partitions.size, schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)
+  end
+
   private
 
   def find_partitioned_columns(table)
@@ -54,4 +69,18 @@ module PartitioningHelpers
         and pg_class.relispartition
     SQL
   end
+
+  def find_partitions(partition, schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)
+    connection.select_rows(<<~SQL)
+      select
+        pg_class.relname
+      from pg_class
+      inner join pg_inherits i on pg_class.oid = inhrelid
+      inner join pg_class parent_class on parent_class.oid = inhparent
+      inner join pg_namespace ON pg_namespace.oid = pg_class.relnamespace
+      where pg_namespace.nspname = '#{schema}'
+        and parent_class.relname = '#{partition}'
+        and pg_class.relispartition
+    SQL
+  end
 end