1 files changed, 19 insertions, 0 deletions
diff --git a/spec/factories/notes.rb b/spec/factories/notes.rb
index fdd1a9a18b2..7c3ba122b5a 100644
--- a/spec/factories/notes.rb
+++ b/spec/factories/notes.rb
@@ -16,6 +16,7 @@ FactoryBot.define do
     factory :note_on_merge_request,      traits: [:on_merge_request]
     factory :note_on_project_snippet,    traits: [:on_project_snippet]
     factory :note_on_personal_snippet,   traits: [:on_personal_snippet]
+    factory :note_on_design,             traits: [:on_design]
     factory :system_note,                traits: [:system]
 
     factory :discussion_note, class: 'DiscussionNote'
@@ -107,6 +108,10 @@ FactoryBot.define do
       end
     end
 
+    factory :diff_note_on_design, parent: :note, traits: [:on_design], class: 'DiffNote' do
+      position { build(:image_diff_position, file: noteable.full_path, diff_refs: noteable.diff_refs) }
+    end
+
     trait :on_commit do
       association :project, :repository
       noteable { nil }
@@ -136,6 +141,20 @@ FactoryBot.define do
       project { nil }
     end
 
+    trait :on_design do
+      transient do
+        issue { association(:issue, project: project) }
+      end
+      noteable { association(:design, :with_file, issue: issue) }
+
+      after(:build) do |note|
+        next if note.project == note.noteable.project
+
+        # note validations require consistency between these two objects
+        note.project = note.noteable.project
+      end
+    end
+
     trait :system do
       system { true }
     end