Add exportable getter/setter to Vulnerabilities::Export model

ee/app/models/vulnerabilities/export.rb

@@ -16,8 +16,6 @@ module Vulnerabilities
       csv: 0
     }
 
-    validates :project, presence: true, unless: :group
-    validates :group, presence: true, unless: :project
     validates :status, presence: true
     validates :format, presence: true
     validates :file, presence: true, if: :finished?
@@ -49,6 +47,22 @@ module Vulnerabilities
       end
     end
 
+    def exportable
+      project || author
+    end
+
+    def exportable=(value)
+      case value
+      when Project
+        self.project = value
+      when User
+        self.project = nil
+        self.author = value
+      else
+        raise "Can not assign #{value.class} as exportable"
+      end
+    end
+
     def completed?
       finished? || failed?
     end

ee/spec/factories/vulnerabilities/exports.rb

@@ -38,5 +38,10 @@ FactoryBot.define do
       project { nil }
       group
     end
+
+    trait :user do
+      project { nil }
+      group { nil }
+    end
   end
 end

ee/spec/models/vulnerabilities/export_spec.rb

@@ -23,16 +23,6 @@ describe Vulnerabilities::Export do
 
       it { is_expected.to validate_presence_of(:file) }
     end
-
-    context 'when the group is not set' do
-      it { is_expected.to validate_presence_of(:project) }
-    end
-
-    context 'when the project is not set' do
-      subject { build(:vulnerability_export, :group) }
-
-      it { is_expected.to validate_presence_of(:group) }
-    end
   end
 
   describe '#status' do
@@ -77,6 +67,54 @@ describe Vulnerabilities::Export do
     end
   end
 
+  describe '#exportable' do
+    subject { vulnerability_export.exportable }
+
+    context 'when the export has project assigned' do
+      let(:project) { build(:project) }
+      let(:vulnerability_export) { build(:vulnerability_export, project: project) }
+
+      it { is_expected.to eql(project) }
+    end
+
+    context 'when the export does not have project assigned' do
+      let(:author) { build(:user) }
+      let(:vulnerability_export) { build(:vulnerability_export, :user, author: author) }
+
+      it { is_expected.to eql(author) }
+    end
+  end
+
+  describe '#exportable=' do
+    let(:vulnerability_export) { build(:vulnerability_export) }
+
+    subject(:set_exportable) { vulnerability_export.exportable = exportable }
+
+    context 'when the exportable is a Project' do
+      let(:exportable) { build(:project) }
+
+      it 'changes the exportable of the export to given project' do
+        expect { set_exportable }.to change { vulnerability_export.exportable }.to(exportable)
+      end
+    end
+
+    context 'when the exportable is a User' do
+      let(:exportable) { build(:user) }
+
+      it 'changes the exportable of the export to given user' do
+        expect { set_exportable }.to change { vulnerability_export.exportable }.to(exportable)
+      end
+    end
+
+    context 'when the exportable is a String' do
+      let(:exportable) { 'Foo' }
+
+      it 'raises an exception' do
+        expect { set_exportable }.to raise_error(RuntimeError)
+      end
+    end
+  end
+
   describe '#completed?' do
     context 'when status is created' do
       subject { build(:vulnerability_export, :created) }