1 files changed, 43 insertions, 3 deletions
diff --git a/spec/support/helpers/partitioning_helpers.rb b/spec/support/helpers/partitioning_helpers.rb
index 98a13915d76..8981fea04d5 100644
--- a/spec/support/helpers/partitioning_helpers.rb
+++ b/spec/support/helpers/partitioning_helpers.rb
@@ -9,13 +9,36 @@ module PartitioningHelpers
   end
 
   def expect_range_partition_of(partition_name, table_name, min_value, max_value)
-    definition = find_partition_definition(partition_name)
+    definition = find_partition_definition(partition_name, schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)
 
     expect(definition).not_to be_nil
     expect(definition['base_table']).to eq(table_name.to_s)
     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
+
+  def expect_hash_partition_of(partition_name, table_name, modulus, remainder)
+    definition = find_partition_definition(partition_name, schema: Gitlab::Database::STATIC_PARTITIONS_SCHEMA)
+
+    expect(definition).not_to be_nil
+    expect(definition['base_table']).to eq(table_name.to_s)
+    expect(definition['condition']).to eq("FOR VALUES WITH (modulus #{modulus}, remainder #{remainder})")
+  end
+
   private
 
   def find_partitioned_columns(table)
@@ -40,7 +63,7 @@ module PartitioningHelpers
     SQL
   end
 
-  def find_partition_definition(partition)
+  def find_partition_definition(partition, schema: )
     connection.select_one(<<~SQL)
       select
         parent_class.relname as base_table,
@@ -48,7 +71,24 @@ module PartitioningHelpers
       from pg_class
       inner join pg_inherits i on pg_class.oid = inhrelid
       inner join pg_class parent_class on parent_class.oid = inhparent
-      where pg_class.relname = '#{partition}' and pg_class.relispartition;
+      inner join pg_namespace ON pg_namespace.oid = pg_class.relnamespace
+      where pg_namespace.nspname = '#{schema}'
+        and pg_class.relname = '#{partition}'
+        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