1 files changed, 11 insertions, 4 deletions
diff --git a/spec/lib/banzai/issuable_extractor_spec.rb b/spec/lib/banzai/issuable_extractor_spec.rb
index b2c869bd066..5bbd98592e7 100644
--- a/spec/lib/banzai/issuable_extractor_spec.rb
+++ b/spec/lib/banzai/issuable_extractor_spec.rb
@@ -7,6 +7,7 @@ RSpec.describe Banzai::IssuableExtractor, feature_category: :team_planning do
   let(:user) { create(:user) }
   let(:extractor) { described_class.new(Banzai::RenderContext.new(project, user)) }
   let(:issue) { create(:issue, project: project) }
+  let(:work_item) { create(:work_item, project: project) }
   let(:merge_request) { create(:merge_request, source_project: project) }
   let(:issue_link) do
     html_to_node(
@@ -14,6 +15,12 @@ RSpec.describe Banzai::IssuableExtractor, feature_category: :team_planning do
     )
   end
 
+  let(:work_item_link) do
+    html_to_node(
+      "<a href='' data-work-item='#{work_item.id}' data-reference-type='work_item' class='gfm'>text</a>"
+    )
+  end
+
   let(:merge_request_link) do
     html_to_node(
       "<a href='' data-merge-request='#{merge_request.id}' data-reference-type='merge_request' class='gfm'>text</a>"
@@ -27,17 +34,17 @@ RSpec.describe Banzai::IssuableExtractor, feature_category: :team_planning do
   end
 
   it 'returns instances of issuables for nodes with references' do
-    result = extractor.extract([issue_link, merge_request_link])
+    result = extractor.extract([issue_link, work_item_link, merge_request_link])
 
-    expect(result).to eq(issue_link => issue, merge_request_link => merge_request)
+    expect(result).to eq(issue_link => issue, work_item_link => work_item, merge_request_link => merge_request)
   end
 
   describe 'caching', :request_store do
     it 'saves records to cache' do
-      extractor.extract([issue_link, merge_request_link])
+      extractor.extract([issue_link, work_item_link, merge_request_link])
 
       second_call_queries = ActiveRecord::QueryRecorder.new do
-        extractor.extract([issue_link, merge_request_link])
+        extractor.extract([issue_link, work_item_link, merge_request_link])
       end.count
 
       expect(second_call_queries).to eq 0