Hide the milestone variable on board when the feature is disabled

And separate EE/CE for board.rb

app/models/board.rb

 class Board < ActiveRecord::Base
+  prepend EE::Board
+
   belongs_to :project
-  belongs_to :milestone
 
   has_many :lists, -> { order(:list_type, :position) }, dependent: :delete_all
 
@@ -13,26 +14,4 @@ class Board < ActiveRecord::Base
   def closed_list
     lists.merge(List.closed).take
   end
-
-  def milestone
-    if milestone_id == Milestone::Upcoming.id
-      Milestone::Upcoming
-    else
-      super
-    end
-  end
-
-  def as_json(options = {})
-    milestone_attrs = options.fetch(:include, {})
-                             .extract!(:milestone)
-                             .dig(:milestone, :only)
-
-    super(options).tap do |json|
-      if milestone.present? && milestone_attrs.present?
-        json[:milestone] = milestone_attrs.each_with_object({}) do |attr, json|
-          json[attr] = milestone.public_send(attr)
-        end
-      end
-    end
-  end
 end

app/models/ee/board.rb

+module EE
+  module Board
+    extend ActiveSupport::Concern
+
+    prepended do
+      belongs_to :milestone
+    end
+
+    def milestone
+      return nil unless project.feature_available?(:issue_board_milestone)
+
+      if milestone_id == ::Milestone::Upcoming.id
+        ::Milestone::Upcoming
+      else
+        super
+      end
+    end
+
+    def as_json(options = {})
+      milestone_attrs = options.fetch(:include, {})
+                          .extract!(:milestone)
+                          .dig(:milestone, :only)
+
+      super(options).tap do |json|
+        if milestone.present? && milestone_attrs.present?
+          json[:milestone] = milestone_attrs.each_with_object({}) do |attr, json|
+            json[attr] = milestone.public_send(attr)
+          end
+        end
+      end
+    end
+  end
+end

spec/models/board_spec.rb

@@ -12,27 +12,4 @@ describe Board do
     it { is_expected.to validate_presence_of(:name) }
     it { is_expected.to validate_presence_of(:project) }
   end
-
-  describe 'milestone' do
-    subject { described_class.new }
-
-    it 'returns Milestone::Upcoming for upcoming milestone id' do
-      subject.milestone_id = Milestone::Upcoming.id
-
-      expect(subject.milestone).to eq Milestone::Upcoming
-    end
-
-    it 'returns milestone for valid milestone id' do
-      milestone = create(:milestone)
-      subject.milestone_id = milestone.id
-
-      expect(subject.milestone).to eq milestone
-    end
-
-    it 'returns nil for invalid milestone id' do
-      subject.milestone_id = -1
-
-      expect(subject.milestone).to be_nil
-    end
-  end
 end

spec/models/ee/board_spec.rb

+require 'spec_helper'
+
+describe Board do
+  describe 'milestone' do
+    subject { build(:board) }
+
+    context 'when the feature is available' do
+      before do
+        stub_licensed_features(issue_board_milestone: true)
+      end
+
+      it 'returns Milestone::Upcoming for upcoming milestone id' do
+        subject.milestone_id = Milestone::Upcoming.id
+
+        expect(subject.milestone).to eq Milestone::Upcoming
+      end
+
+      it 'returns milestone for valid milestone id' do
+        milestone = create(:milestone)
+        subject.milestone_id = milestone.id
+
+        expect(subject.milestone).to eq milestone
+      end
+
+      it 'returns nil for invalid milestone id' do
+        subject.milestone_id = -1
+
+        expect(subject.milestone).to be_nil
+      end
+    end
+
+    it 'returns nil when the feature is not available' do
+      stub_licensed_features(issue_board_milestone: false)
+      milestone = create(:milestone)
+      subject.milestone_id = milestone.id
+
+      expect(subject.milestone).to be_nil
+    end
+  end
+end