1 files changed, 93 insertions, 0 deletions
diff --git a/spec/lib/gitlab/background_migration/backfill_vulnerability_reads_cluster_agent_spec.rb b/spec/lib/gitlab/background_migration/backfill_vulnerability_reads_cluster_agent_spec.rb
new file mode 100644
index 00000000000..79699375a8d
--- /dev/null
+++ b/spec/lib/gitlab/background_migration/backfill_vulnerability_reads_cluster_agent_spec.rb
@@ -0,0 +1,93 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::BackgroundMigration::BackfillVulnerabilityReadsClusterAgent, :migration, schema: 20220525221133 do # rubocop:disable Layout/LineLength
+  let(:migration) do
+    described_class.new(start_id: 1, end_id: 10,
+                        batch_table: table_name, batch_column: batch_column,
+                        sub_batch_size: sub_batch_size, pause_ms: pause_ms,
+                        connection: ApplicationRecord.connection)
+  end
+
+  let(:users_table) { table(:users) }
+  let(:vulnerability_reads_table) { table(:vulnerability_reads) }
+  let(:vulnerability_scanners_table) { table(:vulnerability_scanners) }
+  let(:vulnerabilities_table) { table(:vulnerabilities) }
+  let(:namespaces_table) { table(:namespaces) }
+  let(:projects_table) { table(:projects) }
+  let(:cluster_agents_table) { table(:cluster_agents) }
+
+  let(:table_name) { 'vulnerability_reads' }
+  let(:batch_column) { :id }
+  let(:sub_batch_size) { 1_000 }
+  let(:pause_ms) { 0 }
+
+  subject(:perform_migration) { migration.perform }
+
+  before do
+    users_table.create!(id: 1, name: 'John Doe', email: 'test@example.com', projects_limit: 5)
+
+    namespaces_table.create!(id: 1, name: 'Namespace 1', path: 'namespace-1')
+    namespaces_table.create!(id: 2, name: 'Namespace 2', path: 'namespace-2')
+
+    projects_table.create!(id: 1, namespace_id: 1, name: 'Project 1', path: 'project-1', project_namespace_id: 1)
+    projects_table.create!(id: 2, namespace_id: 2, name: 'Project 2', path: 'project-2', project_namespace_id: 2)
+
+    cluster_agents_table.create!(id: 1, name: 'Agent 1', project_id: 1)
+    cluster_agents_table.create!(id: 2, name: 'Agent 2', project_id: 2)
+
+    vulnerability_scanners_table.create!(id: 1, project_id: 1, external_id: 'starboard', name: 'Starboard')
+    vulnerability_scanners_table.create!(id: 2, project_id: 2, external_id: 'starboard', name: 'Starboard')
+
+    add_vulnerability_read!(1, project_id: 1, cluster_agent_id: 1, report_type: 7)
+    add_vulnerability_read!(3, project_id: 1, cluster_agent_id: 2, report_type: 7)
+    add_vulnerability_read!(5, project_id: 2, cluster_agent_id: 2, report_type: 5)
+    add_vulnerability_read!(7, project_id: 2, cluster_agent_id: 3, report_type: 7)
+    add_vulnerability_read!(9, project_id: 2, cluster_agent_id: 2, report_type: 7)
+    add_vulnerability_read!(10, project_id: 1, cluster_agent_id: 1, report_type: 7)
+    add_vulnerability_read!(11, project_id: 1, cluster_agent_id: 1, report_type: 7)
+  end
+
+  it 'backfills `casted_cluster_agent_id` for the selected records', :aggregate_failures do
+    queries = ActiveRecord::QueryRecorder.new do
+      perform_migration
+    end
+
+    expect(queries.count).to eq(3)
+    expect(vulnerability_reads_table.where.not(casted_cluster_agent_id: nil).count).to eq 3
+    expect(vulnerability_reads_table.where.not(casted_cluster_agent_id: nil).pluck(:id)).to match_array([1, 9, 10])
+  end
+
+  it 'tracks timings of queries' do
+    expect(migration.batch_metrics.timings).to be_empty
+
+    expect { perform_migration }.to change { migration.batch_metrics.timings }
+  end
+
+  private
+
+  def add_vulnerability_read!(id, project_id:, cluster_agent_id:, report_type:)
+    vulnerabilities_table.create!(
+      id: id,
+      project_id: project_id,
+      author_id: 1,
+      title: "Vulnerability #{id}",
+      severity: 5,
+      confidence: 5,
+      report_type: report_type
+    )
+
+    vulnerability_reads_table.create!(
+      id: id,
+      uuid: SecureRandom.uuid,
+      severity: 5,
+      state: 1,
+      vulnerability_id: id,
+      scanner_id: project_id,
+      cluster_agent_id: cluster_agent_id.to_s,
+      project_id: project_id,
+      report_type: report_type
+    )
+  end
+end