4 files changed, 145 insertions, 0 deletions
diff --git a/spec/support/models/ci/partitioning_testing/cascade_check.rb b/spec/support/models/ci/partitioning_testing/cascade_check.rb
new file mode 100644
index 00000000000..f553a47ef4f
--- /dev/null
+++ b/spec/support/models/ci/partitioning_testing/cascade_check.rb
@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module PartitioningTesting
+  module CascadeCheck
+    extend ActiveSupport::Concern
+
+    included do
+      after_create :check_partition_cascade_value
+    end
+
+    def check_partition_cascade_value
+      raise 'Partition value not found' unless partition_scope_value
+
+      return if partition_id == partition_scope_value
+
+      raise "partition_id was expected to equal #{partition_scope_value} but it was #{partition_id}."
+    end
+  end
+end
+
+Ci::Partitionable::Testing::PARTITIONABLE_MODELS.each do |klass|
+  next if klass == 'Ci::Pipeline'
+
+  model = klass.safe_constantize
+
+  model.include(PartitioningTesting::CascadeCheck)
+end
diff --git a/spec/support/models/ci/partitioning_testing/partition_identifiers.rb b/spec/support/models/ci/partitioning_testing/partition_identifiers.rb
new file mode 100644
index 00000000000..aa091095fb6
--- /dev/null
+++ b/spec/support/models/ci/partitioning_testing/partition_identifiers.rb
@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module Ci
+  module PartitioningTesting
+    module PartitionIdentifiers
+      module_function
+
+      def ci_testing_partition_id
+        99999
+      end
+    end
+  end
+end
diff --git a/spec/support/models/ci/partitioning_testing/rspec_hooks.rb b/spec/support/models/ci/partitioning_testing/rspec_hooks.rb
new file mode 100644
index 00000000000..39b15ba8721
--- /dev/null
+++ b/spec/support/models/ci/partitioning_testing/rspec_hooks.rb
@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+RSpec.configure do |config|
+  config.include Ci::PartitioningTesting::PartitionIdentifiers
+
+  config.around(:each, :ci_partitionable) do |example|
+    Ci::PartitioningTesting::SchemaHelpers.with_routing_tables do
+      example.run
+    end
+  end
+
+  config.before(:all) do
+    Ci::PartitioningTesting::SchemaHelpers.setup
+  end
+
+  config.after(:all) do
+    Ci::PartitioningTesting::SchemaHelpers.teardown
+  end
+end
diff --git a/spec/support/models/ci/partitioning_testing/schema_helpers.rb b/spec/support/models/ci/partitioning_testing/schema_helpers.rb
new file mode 100644
index 00000000000..712178710da
--- /dev/null
+++ b/spec/support/models/ci/partitioning_testing/schema_helpers.rb
@@ -0,0 +1,86 @@
+# frozen_string_literal: true
+
+module Ci
+  module PartitioningTesting
+    module SchemaHelpers
+      DEFAULT_PARTITION = 100
+
+      module_function
+
+      def with_routing_tables
+        Ci::BuildMetadata.table_name = :p_ci_builds_metadata
+        yield
+      ensure
+        Ci::BuildMetadata.table_name = :ci_builds_metadata
+      end
+
+      # We're dropping the default values here to ensure that the application code
+      # populates the `partition_id` value and it's not falling back on the
+      # database default one. We should be able to clean this up after
+      # partitioning the tables and substituting the routing table in the model:
+      # https://gitlab.com/gitlab-org/gitlab/-/issues/377822
+      #
+      def setup(connection: Ci::ApplicationRecord.connection)
+        each_partitionable_table do |table_name|
+          change_column_default(table_name, from: DEFAULT_PARTITION, to: nil, connection: connection)
+          change_column_default("p_#{table_name}", from: DEFAULT_PARTITION, to: nil, connection: connection)
+          create_test_partition(table_name, connection: connection)
+        end
+      end
+
+      def teardown(connection: Ci::ApplicationRecord.connection)
+        each_partitionable_table do |table_name|
+          drop_test_partition(table_name, connection: connection)
+          change_column_default(table_name, from: nil, to: DEFAULT_PARTITION, connection: connection)
+          change_column_default("p_#{table_name}", from: nil, to: DEFAULT_PARTITION, connection: connection)
+        end
+      end
+
+      def each_partitionable_table
+        ::Ci::Partitionable::Testing::PARTITIONABLE_MODELS.each do |klass|
+          model = klass.safe_constantize
+          table_name = model.table_name.delete_prefix('p_')
+
+          yield(table_name)
+
+          model.reset_column_information if model.connected?
+        end
+      end
+
+      def change_column_default(table_name, from:, to:, connection:)
+        return unless table_available?(table_name, connection: connection)
+
+        connection.change_column_default(table_name, :partition_id, from: from, to: to)
+      end
+
+      def create_test_partition(table_name, connection:)
+        return unless table_available?("p_#{table_name}", connection: connection)
+
+        drop_test_partition(table_name, connection: connection)
+
+        connection.execute(<<~SQL)
+          CREATE TABLE #{full_partition_name(table_name)}
+            PARTITION OF p_#{table_name}
+            FOR VALUES IN (#{PartitioningTesting::PartitionIdentifiers.ci_testing_partition_id});
+        SQL
+      end
+
+      def drop_test_partition(table_name, connection:)
+        return unless table_available?(table_name, connection: connection)
+
+        connection.execute(<<~SQL)
+          DROP TABLE IF EXISTS #{full_partition_name(table_name)};
+        SQL
+      end
+
+      def table_available?(table_name, connection:)
+        connection.table_exists?(table_name) &&
+          connection.column_exists?(table_name, :partition_id)
+      end
+
+      def full_partition_name(table_name)
+        "#{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}._test_gitlab_#{table_name}_partition"
+      end
+    end
+  end
+end