1 files changed, 136 insertions, 0 deletions
diff --git a/spec/lib/gitlab/background_migration/backfill_has_remediations_of_vulnerability_reads_spec.rb b/spec/lib/gitlab/background_migration/backfill_has_remediations_of_vulnerability_reads_spec.rb
new file mode 100644
index 00000000000..0e7a0210758
--- /dev/null
+++ b/spec/lib/gitlab/background_migration/backfill_has_remediations_of_vulnerability_reads_spec.rb
@@ -0,0 +1,136 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::BackgroundMigration::BackfillHasRemediationsOfVulnerabilityReads,
+  feature_category: :database do
+  let(:namespaces) { table(:namespaces) }
+  let(:projects) { table(:projects) }
+  let(:users) { table(:users) }
+  let(:scanners) { table(:vulnerability_scanners) }
+  let(:vulnerabilities) { table(:vulnerabilities) }
+  let(:vulnerability_reads) { table(:vulnerability_reads) }
+
+  let(:namespace) { namespaces.create!(name: 'user', path: 'user') }
+  let(:project) { projects.create!(namespace_id: namespace.id, project_namespace_id: namespace.id) }
+  let(:user) { users.create!(username: 'john_doe', email: 'johndoe@gitlab.com', projects_limit: 10) }
+  let(:scanner) { scanners.create!(project_id: project.id, external_id: 'external_id', name: 'Test Scanner') }
+
+  let(:vulnerability_1) { create_vulnerability(title: 'vulnerability 1') }
+  let(:vulnerability_2) { create_vulnerability(title: 'vulnerability 2') }
+
+  let!(:vulnerability_read_1) { create_vulnerability_read(vulnerability_id: vulnerability_1.id) }
+  let!(:vulnerability_read_2) { create_vulnerability_read(vulnerability_id: vulnerability_2.id) }
+
+  let(:vulnerability_findings) { table(:vulnerability_occurrences) }
+  let(:vulnerability_findings_remediations) { table(:vulnerability_findings_remediations) }
+  let(:vulnerability_remediations) { table(:vulnerability_remediations) }
+  let(:vulnerability_identifiers) { table(:vulnerability_identifiers) }
+
+  subject(:perform_migration) do
+    described_class.new(
+      start_id: vulnerability_reads.first.vulnerability_id,
+      end_id: vulnerability_reads.last.vulnerability_id,
+      batch_table: :vulnerability_reads,
+      batch_column: :vulnerability_id,
+      sub_batch_size: vulnerability_reads.count,
+      pause_ms: 0,
+      connection: ActiveRecord::Base.connection
+    ).perform
+  end
+
+  it 'updates vulnerability_reads records which has remediations' do
+    vuln_remediation = create_remediation
+    vuln_finding = create_finding(vulnerability_id: vulnerability_1.id)
+    vulnerability_findings_remediations.create!(
+      vulnerability_occurrence_id: vuln_finding.id,
+      vulnerability_remediation_id: vuln_remediation.id
+    )
+
+    expect { perform_migration }.to change { vulnerability_read_1.reload.has_remediations }.from(false).to(true)
+                                .and not_change { vulnerability_read_2.reload.has_remediations }.from(false)
+  end
+
+  it 'does not modify has_remediations of vulnerabilities which do not have remediations' do
+    expect { perform_migration }.to not_change { vulnerability_read_1.reload.has_remediations }.from(false)
+                                .and not_change { vulnerability_read_2.reload.has_remediations }.from(false)
+  end
+
+  private
+
+  def create_vulnerability(overrides = {})
+    attrs = {
+      project_id: project.id,
+      author_id: user.id,
+      title: 'test',
+      severity: 1,
+      confidence: 1,
+      report_type: 1
+    }.merge(overrides)
+
+    vulnerabilities.create!(attrs)
+  end
+
+  def create_vulnerability_read(overrides = {})
+    attrs = {
+      project_id: project.id,
+      vulnerability_id: 1,
+      scanner_id: scanner.id,
+      severity: 1,
+      report_type: 1,
+      state: 1,
+      uuid: SecureRandom.uuid
+    }.merge(overrides)
+
+    vulnerability_reads.create!(attrs)
+  end
+
+  def create_finding(overrides = {})
+    attrs = {
+      project_id: project.id,
+      scanner_id: scanner.id,
+      severity: 5, # medium
+      confidence: 2, # unknown,
+      report_type: 99, # generic
+      primary_identifier_id: create_identifier.id,
+      project_fingerprint: SecureRandom.hex(20),
+      location_fingerprint: SecureRandom.hex(20),
+      uuid: SecureRandom.uuid,
+      name: "CVE-2018-1234",
+      raw_metadata: "{}",
+      metadata_version: "test:1.0"
+    }.merge(overrides)
+
+    vulnerability_findings.create!(attrs)
+  end
+
+  def create_remediation(overrides = {})
+    remediation_hash = { summary: 'summary', diff: "ZGlmZiAtLWdp" }
+
+    attrs = {
+      project_id: project.id,
+      summary: remediation_hash[:summary],
+      checksum: checksum(remediation_hash[:diff]),
+      file: Tempfile.new.path
+    }.merge(overrides)
+
+    vulnerability_remediations.create!(attrs)
+  end
+
+  def create_identifier(overrides = {})
+    attrs = {
+      project_id: project.id,
+      external_id: "CVE-2018-1234",
+      external_type: "CVE",
+      name: "CVE-2018-1234",
+      fingerprint: SecureRandom.hex(20)
+    }.merge(overrides)
+
+    vulnerability_identifiers.create!(attrs)
+  end
+
+  def checksum(value)
+    sha = Digest::SHA256.hexdigest(value)
+    Gitlab::Database::ShaAttribute.new.serialize(sha)
+  end
+end