3 files changed, 31 insertions, 35 deletions
diff --git a/spec/models/work_items/type_spec.rb b/spec/models/work_items/type_spec.rb
index e41df7f0f61..6685720778a 100644
--- a/spec/models/work_items/type_spec.rb
+++ b/spec/models/work_items/type_spec.rb
@@ -35,10 +35,10 @@ RSpec.describe WorkItems::Type do
       it 'deletes type but not unrelated issues' do
         type = create(:work_item_type)
 
-        expect(WorkItems::Type.count).to eq(6)
+        expect(WorkItems::Type.count).to eq(8)
 
         expect { type.destroy! }.not_to change(Issue, :count)
-        expect(WorkItems::Type.count).to eq(5)
+        expect(WorkItems::Type.count).to eq(7)
       end
     end
 
@@ -69,7 +69,8 @@ RSpec.describe WorkItems::Type do
         ::WorkItems::Widgets::Hierarchy,
         ::WorkItems::Widgets::Labels,
         ::WorkItems::Widgets::Assignees,
-        ::WorkItems::Widgets::StartAndDueDate
+        ::WorkItems::Widgets::StartAndDueDate,
+        ::WorkItems::Widgets::Milestone
       )
     end
   end
diff --git a/spec/models/work_items/widgets/hierarchy_spec.rb b/spec/models/work_items/widgets/hierarchy_spec.rb
index cd528772710..c847f2694fe 100644
--- a/spec/models/work_items/widgets/hierarchy_spec.rb
+++ b/spec/models/work_items/widgets/hierarchy_spec.rb
@@ -26,22 +26,6 @@ RSpec.describe WorkItems::Widgets::Hierarchy do
     subject { described_class.new(parent_link.work_item).parent }
 
     it { is_expected.to eq(parent_link.work_item_parent) }
-
-    context 'when work_items flag is disabled' do
-      before do
-        stub_feature_flags(work_items: false)
-      end
-
-      it { is_expected.to be_nil }
-    end
-
-    context 'when work_items flag is enabled for the parent group' do
-      before do
-        stub_feature_flags(work_items: group)
-      end
-
-      it { is_expected.to eq(parent_link.work_item_parent) }
-    end
   end
 
   describe '#children' do
@@ -51,21 +35,5 @@ RSpec.describe WorkItems::Widgets::Hierarchy do
     subject { described_class.new(work_item_parent).children }
 
     it { is_expected.to contain_exactly(parent_link1.work_item, parent_link2.work_item) }
-
-    context 'when work_items flag is disabled' do
-      before do
-        stub_feature_flags(work_items: false)
-      end
-
-      it { is_expected.to be_empty }
-    end
-
-    context 'when work_items flag is enabled for the parent group' do
-      before do
-        stub_feature_flags(work_items: group)
-      end
-
-      it { is_expected.to contain_exactly(parent_link1.work_item, parent_link2.work_item) }
-    end
   end
 end
diff --git a/spec/models/work_items/widgets/milestone_spec.rb b/spec/models/work_items/widgets/milestone_spec.rb
new file mode 100644
index 00000000000..7b2d661df29
--- /dev/null
+++ b/spec/models/work_items/widgets/milestone_spec.rb
@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe WorkItems::Widgets::Milestone do
+  let_it_be(:project) { create(:project) }
+  let_it_be(:milestone) { create(:milestone, project: project) }
+  let_it_be(:work_item) { create(:work_item, :issue, project: project, milestone: milestone) }
+
+  describe '.type' do
+    subject { described_class.type }
+
+    it { is_expected.to eq(:milestone) }
+  end
+
+  describe '#type' do
+    subject { described_class.new(work_item).type }
+
+    it { is_expected.to eq(:milestone) }
+  end
+
+  describe '#milestone' do
+    subject { described_class.new(work_item).milestone }
+
+    it { is_expected.to eq(work_item.milestone) }
+  end
+end