Add latest changes from gitlab-org/gitlab@master

38 files changed, 358 insertions, 81 deletions

diff --git a/spec/factories/ci/stages.rb b/spec/factories/ci/stages.rb
index 65fca3fc314..41297b01f92 100644
--- a/spec/factories/ci/stages.rb
+++ b/spec/factories/ci/stages.rb
@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 FactoryBot.define do
-  factory :ci_stage_entity, class: 'Ci::Stage' do
+  factory :ci_stage, class: 'Ci::Stage' do
     project factory: :project
     pipeline factory: :ci_empty_pipeline
 
diff --git a/spec/features/incidents/incident_timeline_events_spec.rb b/spec/features/incidents/incident_timeline_events_spec.rb
new file mode 100644
index 00000000000..4424af6a489
--- /dev/null
+++ b/spec/features/incidents/incident_timeline_events_spec.rb
@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe 'Incident timeline events', :js do
+  let_it_be(:project) { create(:project) }
+  let_it_be(:developer) { create(:user) }
+  let_it_be(:incident) { create(:incident, project: project) }
+
+  before_all do
+    project.add_developer(developer)
+  end
+
+  before do
+    stub_feature_flags(incident_timeline: true)
+    sign_in(developer)
+
+    visit project_issues_incident_path(project, incident)
+    wait_for_requests
+    click_link 'Timeline'
+  end
+
+  context 'when add event is clicked' do
+    it 'submits event data when save is clicked' do
+      click_button 'Add new timeline event'
+
+      expect(page).to have_selector('.common-note-form')
+
+      fill_in 'Description', with: 'Event note goes here'
+      fill_in 'timeline-input-hours', with: '07'
+      fill_in 'timeline-input-minutes', with: '25'
+
+      click_button 'Save'
+
+      expect(page).to have_selector('.incident-timeline-events')
+
+      page.within '.timeline-event-note' do
+        expect(page).to have_content('Event note goes here')
+        expect(page).to have_content('07:25')
+      end
+    end
+  end
+end
diff --git a/spec/frontend/issues/show/components/incidents/mock_data.js b/spec/frontend/issues/show/components/incidents/mock_data.js
index b5346a6089a..0c58a44a6bc 100644
--- a/spec/frontend/issues/show/components/incidents/mock_data.js
+++ b/spec/frontend/issues/show/components/incidents/mock_data.js
@@ -70,3 +70,17 @@ export const timelineEventsQueryEmptyResponse = {
     },
   },
 };
+
+export const timelineEventsCreateEventResponse = {
+  timelineEvent: {
+    ...mockEvents[0],
+  },
+  errors: [],
+};
+
+export const timelineEventsCreateEventError = {
+  timelineEvent: {
+    ...mockEvents[0],
+  },
+  errors: ['Error creating timeline event'],
+};
diff --git a/spec/frontend/issues/show/components/incidents/timeline_events_form_spec.js b/spec/frontend/issues/show/components/incidents/timeline_events_form_spec.js
new file mode 100644
index 00000000000..2b7f73c0e1c
--- /dev/null
+++ b/spec/frontend/issues/show/components/incidents/timeline_events_form_spec.js
@@ -0,0 +1,141 @@
+import VueApollo from 'vue-apollo';
+import Vue, { nextTick } from 'vue';
+import { GlDatepicker } from '@gitlab/ui';
+import { shallowMountExtended, mountExtended } from 'helpers/vue_test_utils_helper';
+import waitForPromises from 'helpers/wait_for_promises';
+import IncidentTimelineEventForm from '~/issues/show/components/incidents/timeline_events_form.vue';
+import createTimelineEventMutation from '~/issues/show/components/incidents/graphql/queries/create_timeline_event.mutation.graphql';
+import createMockApollo from 'helpers/mock_apollo_helper';
+import { createAlert } from '~/flash';
+import { timelineEventsCreateEventResponse, timelineEventsCreateEventError } from './mock_data';
+
+Vue.use(VueApollo);
+
+jest.mock('~/flash');
+
+const addEventResponse = jest.fn().mockResolvedValue(timelineEventsCreateEventResponse);
+
+function createMockApolloProvider(response = addEventResponse) {
+  const requestHandlers = [[createTimelineEventMutation, response]];
+  return createMockApollo(requestHandlers);
+}
+
+describe('Timeline events form', () => {
+  let wrapper;
+
+  const mountComponent = ({ mockApollo, mountMethod = shallowMountExtended, stubs }) => {
+    wrapper = mountMethod(IncidentTimelineEventForm, {
+      propsData: {
+        hasTimelineEvents: true,
+      },
+      provide: {
+        fullPath: 'group/project',
+        issuableId: '1',
+      },
+      apolloProvider: mockApollo,
+      stubs,
+    });
+  };
+
+  afterEach(() => {
+    if (wrapper) {
+      wrapper.destroy();
+    }
+  });
+
+  const findSubmitButton = () => wrapper.findByText('Save');
+  const findSubmitAndAddButton = () => wrapper.findByText('Save and add another event');
+  const findCancelButton = () => wrapper.findByText('Cancel');
+  const findDatePicker = () => wrapper.findComponent(GlDatepicker);
+  const findHourInput = () => wrapper.findByTestId('input-hours');
+  const findMinuteInput = () => wrapper.findByTestId('input-minutes');
+
+  const submitForm = async () => {
+    findSubmitButton().trigger('click');
+    await waitForPromises();
+  };
+  const submitFormAndAddAnother = async () => {
+    findSubmitAndAddButton().trigger('click');
+    await waitForPromises();
+  };
+  const cancelForm = async () => {
+    findCancelButton().trigger('click');
+    await waitForPromises();
+  };
+
+  describe('form button behaviour', () => {
+    const closeFormEvent = { 'hide-incident-timeline-event-form': [[]] };
+    beforeEach(() => {
+      mountComponent({ mockApollo: createMockApolloProvider(), mountMethod: mountExtended });
+    });
+
+    it('should close the form on submit', async () => {
+      await submitForm();
+      expect(wrapper.emitted()).toEqual(closeFormEvent);
+    });
+
+    it('should not close the form on "submit and add another"', async () => {
+      await submitFormAndAddAnother();
+      expect(wrapper.emitted()).toEqual({});
+    });
+
+    it('should close the form on cancel', async () => {
+      await cancelForm();
+      expect(wrapper.emitted()).toEqual(closeFormEvent);
+    });
+  });
+
+  describe('addTimelineEventQuery', () => {
+    const expectedData = {
+      input: {
+        incidentId: 'gid://gitlab/Issue/1',
+        note: '',
+        occurredAt: '2020-07-06T00:00:00.000Z',
+      },
+    };
+
+    let mockApollo;
+
+    beforeEach(() => {
+      mockApollo = createMockApolloProvider();
+      mountComponent({ mockApollo, mountMethod: mountExtended });
+    });
+
+    it('should call the mutation with the right variables', async () => {
+      await submitForm();
+
+      expect(addEventResponse).toHaveBeenCalledWith(expectedData);
+    });
+
+    it('should call the mutation with user selected variables', async () => {
+      const expectedUserSelectedData = {
+        input: {
+          ...expectedData.input,
+          occurredAt: '2021-08-12T05:45:00.000Z',
+        },
+      };
+
+      findDatePicker().vm.$emit('input', new Date('2021-08-12'));
+      findHourInput().vm.$emit('input', 5);
+      findMinuteInput().vm.$emit('input', 45);
+
+      await nextTick();
+      await submitForm();
+
+      expect(addEventResponse).toHaveBeenCalledWith(expectedUserSelectedData);
+    });
+  });
+
+  describe('error handling', () => {
+    const mockApollo = createMockApolloProvider(timelineEventsCreateEventError);
+    beforeEach(() => {
+      mountComponent({ mockApollo, mountMethod: mountExtended });
+    });
+
+    it('should show an error when submission fails', async () => {
+      await submitForm();
+
+      expect(createAlert).toHaveBeenCalled();
+    });
+  });
+});
diff --git a/spec/frontend/issues/show/components/incidents/timeline_events_tab_spec.js b/spec/frontend/issues/show/components/incidents/timeline_events_tab_spec.js
index cf81f4cdf66..993df59c5fa 100644
--- a/spec/frontend/issues/show/components/incidents/timeline_events_tab_spec.js
+++ b/spec/frontend/issues/show/components/incidents/timeline_events_tab_spec.js
@@ -1,13 +1,15 @@
-import { GlEmptyState, GlLoadingIcon } from '@gitlab/ui';
+import { GlLoadingIcon, GlEmptyState } from '@gitlab/ui';
 import VueApollo from 'vue-apollo';
 import Vue from 'vue';
-import { shallowMountExtended } from 'helpers/vue_test_utils_helper';
+import { shallowMountExtended, mountExtended } from 'helpers/vue_test_utils_helper';
 import waitForPromises from 'helpers/wait_for_promises';
 import TimelineEventsTab from '~/issues/show/components/incidents/timeline_events_tab.vue';
 import IncidentTimelineEventsList from '~/issues/show/components/incidents/timeline_events_list.vue';
+import IncidentTimelineEventForm from '~/issues/show/components/incidents/timeline_events_form.vue';
 import timelineEventsQuery from '~/issues/show/components/incidents/graphql/queries/get_timeline_events.query.graphql';
 import createMockApollo from 'helpers/mock_apollo_helper';
 import { createAlert } from '~/flash';
+import { timelineTabI18n } from '~/issues/show/components/incidents/constants';
 import { timelineEventsQueryListResponse, timelineEventsQueryEmptyResponse } from './mock_data';
 
 Vue.use(VueApollo);
@@ -28,14 +30,17 @@ describe('TimelineEventsTab', () => {
   let wrapper;
 
   const mountComponent = (options = {}) => {
-    const { mockApollo, mountMethod = shallowMountExtended } = options;
+    const { mockApollo, mountMethod = shallowMountExtended, stubs, provide } = options;
 
     wrapper = mountMethod(TimelineEventsTab, {
       provide: {
         fullPath: 'group/project',
         issuableId: '1',
+        canUpdate: true,
+        ...provide,
       },
       apolloProvider: mockApollo,
+      stubs,
     });
   };
 
@@ -48,6 +53,8 @@ describe('TimelineEventsTab', () => {
   const findLoadingSpinner = () => wrapper.findComponent(GlLoadingIcon);
   const findEmptyState = () => wrapper.findComponent(GlEmptyState);
   const findTimelineEventsList = () => wrapper.findComponent(IncidentTimelineEventsList);
+  const findTimelineEventForm = () => wrapper.findComponent(IncidentTimelineEventForm);
+  const findAddEventButton = () => wrapper.findByText(timelineTabI18n.addEventButton);
 
   describe('Timeline events tab', () => {
     describe('empty state', () => {
@@ -102,4 +109,50 @@ describe('TimelineEventsTab', () => {
       expect(findTimelineEventsList().props('timelineEvents')).toHaveLength(3);
     });
   });
+
+  describe('add new event form', () => {
+    beforeEach(async () => {
+      mountComponent({
+        mockApollo: createMockApolloProvider(emptyResponse),
+        mountMethod: mountExtended,
+        stubs: {
+          'incident-timeline-events-list': true,
+          'gl-tab': true,
+        },
+      });
+      await waitForPromises();
+    });
+
+    it('should show a button when user can update', () => {
+      expect(findAddEventButton().exists()).toBe(true);
+    });
+
+    it('should not show a button when user cannot update', () => {
+      mountComponent({
+        mockApollo: createMockApolloProvider(emptyResponse),
+        provide: { canUpdate: false },
+      });
+
+      expect(findAddEventButton().exists()).toBe(false);
+    });
+
+    it('should not show a form by default', () => {
+      expect(findTimelineEventForm().isVisible()).toBe(false);
+    });
+
+    it('should show a form when button is clicked', async () => {
+      await findAddEventButton().trigger('click');
+
+      expect(findTimelineEventForm().isVisible()).toBe(true);
+    });
+
+    it('should hide the form when the hide event is emitted', async () => {
+      // open the form
+      await findAddEventButton().trigger('click');
+
+      await findTimelineEventForm().vm.$emit('hide-incident-timeline-event-form');
+
+      expect(findTimelineEventForm().isVisible()).toBe(false);
+    });
+  });
 });
diff --git a/spec/frontend/issues/show/components/incidents/utils_spec.js b/spec/frontend/issues/show/components/incidents/utils_spec.js
index ff148702696..0da0114c654 100644
--- a/spec/frontend/issues/show/components/incidents/utils_spec.js
+++ b/spec/frontend/issues/show/components/incidents/utils_spec.js
@@ -1,4 +1,9 @@
-import { displayAndLogError, getEventIcon } from '~/issues/show/components/incidents/utils';
+import timezoneMock from 'timezone-mock';
+import {
+  displayAndLogError,
+  getEventIcon,
+  getUtcShiftedDateNow,
+} from '~/issues/show/components/incidents/utils';
 import { createAlert } from '~/flash';
 
 jest.mock('~/flash');
@@ -28,4 +33,22 @@ describe('incident utils', () => {
       expect(getEventIcon('non-existent-icon-name')).toBe('comment');
     });
   });
+
+  describe('getUtcShiftedDateNow', () => {
+    beforeEach(() => {
+      timezoneMock.register('US/Pacific');
+    });
+
+    afterEach(() => {
+      timezoneMock.unregister();
+    });
+
+    it('should shift the date by the timezone offset', () => {
+      const date = new Date();
+
+      const shiftedDate = getUtcShiftedDateNow();
+
+      expect(shiftedDate > date).toBe(true);
+    });
+  });
 });
diff --git a/spec/graphql/types/ci/detailed_status_type_spec.rb b/spec/graphql/types/ci/detailed_status_type_spec.rb
index 0c05227aec2..686461cb9a5 100644
--- a/spec/graphql/types/ci/detailed_status_type_spec.rb
+++ b/spec/graphql/types/ci/detailed_status_type_spec.rb
@@ -13,7 +13,7 @@ RSpec.describe Types::Ci::DetailedStatusType do
                                                    :label, :text, :tooltip, :action)
   end
 
-  let_it_be(:stage) { create(:ci_stage_entity, status: :skipped) }
+  let_it_be(:stage) { create(:ci_stage, status: :skipped) }
 
   describe 'id field' do
     it 'correctly renders the field' do
diff --git a/spec/graphql/types/ci/status_action_type_spec.rb b/spec/graphql/types/ci/status_action_type_spec.rb
index 4c467bf240e..3f4b52610ae 100644
--- a/spec/graphql/types/ci/status_action_type_spec.rb
+++ b/spec/graphql/types/ci/status_action_type_spec.rb
@@ -22,7 +22,7 @@ RSpec.describe Types::Ci::StatusActionType do
 
   describe 'id field' do
     it 'correctly renders the field' do
-      stage = build(:ci_stage_entity, status: :skipped)
+      stage = build(:ci_stage, status: :skipped)
       status = stage.detailed_status(stage.pipeline.user)
 
       expected_id = "#{stage.class.name}-#{status.id}"
diff --git a/spec/lib/gitlab/ci/pipeline/chain/create_deployments_spec.rb b/spec/lib/gitlab/ci/pipeline/chain/create_deployments_spec.rb
index 375841ce236..cbf92f8fa83 100644
--- a/spec/lib/gitlab/ci/pipeline/chain/create_deployments_spec.rb
+++ b/spec/lib/gitlab/ci/pipeline/chain/create_deployments_spec.rb
@@ -6,7 +6,7 @@ RSpec.describe Gitlab::Ci::Pipeline::Chain::CreateDeployments do
   let_it_be(:project) { create(:project, :repository) }
   let_it_be(:user) { create(:user) }
 
-  let(:stage) { build(:ci_stage_entity, project: project, statuses: [job]) }
+  let(:stage) { build(:ci_stage, project: project, statuses: [job]) }
   let(:pipeline) { create(:ci_pipeline, project: project, stages: [stage]) }
 
   let(:command) do
diff --git a/spec/lib/gitlab/ci/pipeline/chain/create_spec.rb b/spec/lib/gitlab/ci/pipeline/chain/create_spec.rb
index 7b97dde3808..eba0db0adfb 100644
--- a/spec/lib/gitlab/ci/pipeline/chain/create_spec.rb
+++ b/spec/lib/gitlab/ci/pipeline/chain/create_spec.rb
@@ -59,7 +59,7 @@ RSpec.describe Gitlab::Ci::Pipeline::Chain::Create do
 
   context 'tags persistence' do
     let(:stage) do
-      build(:ci_stage_entity, pipeline: pipeline, project: project)
+      build(:ci_stage, pipeline: pipeline, project: project)
     end
 
     let(:job) do
diff --git a/spec/lib/gitlab/ci/pipeline/chain/ensure_environments_spec.rb b/spec/lib/gitlab/ci/pipeline/chain/ensure_environments_spec.rb
index 6a7d9b58a05..e07a3ca9033 100644
--- a/spec/lib/gitlab/ci/pipeline/chain/ensure_environments_spec.rb
+++ b/spec/lib/gitlab/ci/pipeline/chain/ensure_environments_spec.rb
@@ -5,7 +5,7 @@ require 'spec_helper'
 RSpec.describe Gitlab::Ci::Pipeline::Chain::EnsureEnvironments do
   let(:project) { create(:project) }
   let(:user) { create(:user) }
-  let(:stage) { build(:ci_stage_entity, project: project, statuses: [job]) }
+  let(:stage) { build(:ci_stage, project: project, statuses: [job]) }
   let(:pipeline) { build(:ci_pipeline, project: project, stages: [stage]) }
 
   let(:command) do
diff --git a/spec/lib/gitlab/ci/pipeline/chain/ensure_resource_groups_spec.rb b/spec/lib/gitlab/ci/pipeline/chain/ensure_resource_groups_spec.rb
index 571455d6279..f14dd70a753 100644
--- a/spec/lib/gitlab/ci/pipeline/chain/ensure_resource_groups_spec.rb
+++ b/spec/lib/gitlab/ci/pipeline/chain/ensure_resource_groups_spec.rb
@@ -5,7 +5,7 @@ require 'spec_helper'
 RSpec.describe Gitlab::Ci::Pipeline::Chain::EnsureResourceGroups do
   let(:project)  { create(:project) }
   let(:user)     { create(:user) }
-  let(:stage) { build(:ci_stage_entity, project: project, statuses: [job]) }
+  let(:stage) { build(:ci_stage, project: project, statuses: [job]) }
   let(:pipeline) { build(:ci_pipeline, project: project, stages: [stage]) }
   let!(:environment) { create(:environment, name: 'production', project: project) }
 
diff --git a/spec/lib/gitlab/ci/status/stage/common_spec.rb b/spec/lib/gitlab/ci/status/stage/common_spec.rb
index a7e0ebf9158..bbd2ce6c83b 100644
--- a/spec/lib/gitlab/ci/status/stage/common_spec.rb
+++ b/spec/lib/gitlab/ci/status/stage/common_spec.rb
@@ -8,7 +8,7 @@ RSpec.describe Gitlab::Ci::Status::Stage::Common do
   let(:pipeline) { create(:ci_empty_pipeline, project: project) }
 
   let(:stage) do
-    build(:ci_stage_entity, pipeline: pipeline, name: 'test')
+    build(:ci_stage, pipeline: pipeline, name: 'test')
   end
 
   subject do
diff --git a/spec/lib/gitlab/ci/status/stage/factory_spec.rb b/spec/lib/gitlab/ci/status/stage/factory_spec.rb
index 348e0a1eb68..35d44281072 100644
--- a/spec/lib/gitlab/ci/status/stage/factory_spec.rb
+++ b/spec/lib/gitlab/ci/status/stage/factory_spec.rb
@@ -7,7 +7,7 @@ RSpec.describe Gitlab::Ci::Status::Stage::Factory do
   let(:project) { create(:project) }
   let(:pipeline) { create(:ci_empty_pipeline, project: project) }
 
-  let(:stage) { create(:ci_stage_entity, pipeline: pipeline) }
+  let(:stage) { create(:ci_stage, pipeline: pipeline) }
 
   subject do
     described_class.new(stage, user)
@@ -24,7 +24,7 @@ RSpec.describe Gitlab::Ci::Status::Stage::Factory do
   context 'when stage has a core status' do
     (Ci::HasStatus::AVAILABLE_STATUSES - %w(manual skipped scheduled)).each do |core_status|
       context "when core status is #{core_status}" do
-        let(:stage) { create(:ci_stage_entity, pipeline: pipeline, status: core_status) }
+        let(:stage) { create(:ci_stage, pipeline: pipeline, status: core_status) }
 
         it "fabricates a core status #{core_status}" do
           expect(status).to be_a(
@@ -42,7 +42,7 @@ RSpec.describe Gitlab::Ci::Status::Stage::Factory do
 
   context 'when stage has warnings' do
     let(:stage) do
-      create(:ci_stage_entity, status: :success, pipeline: pipeline)
+      create(:ci_stage, status: :success, pipeline: pipeline)
     end
 
     before do
@@ -64,7 +64,7 @@ RSpec.describe Gitlab::Ci::Status::Stage::Factory do
   context 'when stage has manual builds' do
     (Ci::HasStatus::BLOCKED_STATUS + ['skipped']).each do |core_status|
       context "when status is #{core_status}" do
-        let(:stage) { create(:ci_stage_entity, pipeline: pipeline, status: core_status) }
+        let(:stage) { create(:ci_stage, pipeline: pipeline, status: core_status) }
 
         it 'fabricates a play manual status' do
           expect(status).to be_a(Gitlab::Ci::Status::Stage::PlayManual)
diff --git a/spec/lib/gitlab/ci/status/stage/play_manual_spec.rb b/spec/lib/gitlab/ci/status/stage/play_manual_spec.rb
index 25b79ff2099..9fdaddc083e 100644
--- a/spec/lib/gitlab/ci/status/stage/play_manual_spec.rb
+++ b/spec/lib/gitlab/ci/status/stage/play_manual_spec.rb
@@ -25,7 +25,7 @@ RSpec.describe Gitlab::Ci::Status::Stage::PlayManual do
   end
 
   describe '#action_path' do
-    let(:stage) { create(:ci_stage_entity, status: 'manual') }
+    let(:stage) { create(:ci_stage, status: 'manual') }
     let(:pipeline) { stage.pipeline }
     let(:play_manual) { stage.detailed_status(create(:user)) }
 
@@ -46,25 +46,25 @@ RSpec.describe Gitlab::Ci::Status::Stage::PlayManual do
     subject { described_class.matches?(stage, user) }
 
     context 'when stage is skipped' do
-      let(:stage) { create(:ci_stage_entity, status: :skipped) }
+      let(:stage) { create(:ci_stage, status: :skipped) }
 
       it { is_expected.to be_truthy }
     end
 
     context 'when stage is manual' do
-      let(:stage) { create(:ci_stage_entity, status: :manual) }
+      let(:stage) { create(:ci_stage, status: :manual) }
 
       it { is_expected.to be_truthy }
     end
 
     context 'when stage is scheduled' do
-      let(:stage) { create(:ci_stage_entity, status: :scheduled) }
+      let(:stage) { create(:ci_stage, status: :scheduled) }
 
       it { is_expected.to be_truthy }
     end
 
     context 'when stage is success' do
-      let(:stage) { create(:ci_stage_entity, status: :success) }
+      let(:stage) { create(:ci_stage, status: :success) }
 
       context 'and does not have manual builds' do
         it { is_expected.to be_falsy }
diff --git a/spec/lib/gitlab/database/loose_foreign_keys_spec.rb b/spec/lib/gitlab/database/loose_foreign_keys_spec.rb
index ed11699e494..87a3e0f81e4 100644
--- a/spec/lib/gitlab/database/loose_foreign_keys_spec.rb
+++ b/spec/lib/gitlab/database/loose_foreign_keys_spec.rb
@@ -63,19 +63,22 @@ RSpec.describe Gitlab::Database::LooseForeignKeys do
         Gitlab::Database.schemas_to_base_models.fetch(parent_table_schema)
       end
 
-      it 'all `to_table` tables are present' do
+      it 'all `to_table` tables are present', :aggregate_failures do
         definitions.each do |definition|
           base_models_for(definition.to_table).each do |model|
-            expect(model.connection).to be_table_exist(definition.to_table)
+            expect(model.connection).to be_table_exist(definition.to_table),
+              "Table #{definition.from_table} does not exist"
           end
         end
       end
 
-      it 'all `from_table` tables are present' do
+      it 'all `from_table` tables are present', :aggregate_failures do
         definitions.each do |definition|
           base_models_for(definition.from_table).each do |model|
-            expect(model.connection).to be_table_exist(definition.from_table)
-            expect(model.connection).to be_column_exist(definition.from_table, definition.column)
+            expect(model.connection).to be_table_exist(definition.from_table),
+              "Table #{definition.from_table} does not exist"
+            expect(model.connection).to be_column_exist(definition.from_table, definition.column),
+              "Column #{definition.column} in #{definition.from_table} does not exist"
           end
         end
       end
diff --git a/spec/models/ci/group_spec.rb b/spec/models/ci/group_spec.rb
index 6c96e659a34..4900bc792af 100644
--- a/spec/models/ci/group_spec.rb
+++ b/spec/models/ci/group_spec.rb
@@ -70,7 +70,7 @@ RSpec.describe Ci::Group do
 
   describe '.fabricate' do
     let(:pipeline) { create(:ci_empty_pipeline) }
-    let(:stage) { create(:ci_stage_entity, pipeline: pipeline) }
+    let(:stage) { create(:ci_stage, pipeline: pipeline) }
 
     before do
       create_build(:ci_build, name: 'rspec 0 2')
diff --git a/spec/models/ci/pipeline_spec.rb b/spec/models/ci/pipeline_spec.rb
index 15cb5b42248..e8973de5fd5 100644
--- a/spec/models/ci/pipeline_spec.rb
+++ b/spec/models/ci/pipeline_spec.rb
@@ -1350,7 +1350,7 @@ RSpec.describe Ci::Pipeline, :mailer, factory_default: :keep do
       let(:pipeline) { build(:ci_empty_pipeline, :created) }
 
       before do
-        create(:ci_stage_entity, project: project,
+        create(:ci_stage, project: project,
                                  pipeline: pipeline,
                                  position: 4,
                                  name: 'deploy')
@@ -1367,12 +1367,12 @@ RSpec.describe Ci::Pipeline, :mailer, factory_default: :keep do
                           stage_idx: 2,
                           name: 'build')
 
-        create(:ci_stage_entity, project: project,
+        create(:ci_stage, project: project,
                                  pipeline: pipeline,
                                  position: 1,
                                  name: 'sanity')
 
-        create(:ci_stage_entity, project: project,
+        create(:ci_stage, project: project,
                                  pipeline: pipeline,
                                  position: 5,
                                  name: 'cleanup')
@@ -4258,7 +4258,7 @@ RSpec.describe Ci::Pipeline, :mailer, factory_default: :keep do
     let(:stage_name) { 'test' }
 
     let(:stage) do
-      create(:ci_stage_entity,
+      create(:ci_stage,
              pipeline: pipeline,
              project: pipeline.project,
              name: 'test')
diff --git a/spec/models/ci/processable_spec.rb b/spec/models/ci/processable_spec.rb
index 743736236c2..789ae3a2ccc 100644
--- a/spec/models/ci/processable_spec.rb
+++ b/spec/models/ci/processable_spec.rb
@@ -24,7 +24,7 @@ RSpec.describe Ci::Processable do
       new_proc
     end
 
-    let_it_be(:stage) { create(:ci_stage_entity, project: project, pipeline: pipeline, name: 'test') }
+    let_it_be(:stage) { create(:ci_stage, project: project, pipeline: pipeline, name: 'test') }
 
     shared_context 'processable bridge' do
       let_it_be(:downstream_project) { create(:project, :repository) }
diff --git a/spec/models/ci/stage_spec.rb b/spec/models/ci/stage_spec.rb
index b91348eb408..d55a8509a98 100644
--- a/spec/models/ci/stage_spec.rb
+++ b/spec/models/ci/stage_spec.rb
@@ -5,7 +5,7 @@ require 'spec_helper'
 RSpec.describe Ci::Stage, :models do
   let_it_be(:pipeline) { create(:ci_empty_pipeline) }
 
-  let(:stage) { create(:ci_stage_entity, pipeline: pipeline, project: pipeline.project) }
+  let(:stage) { create(:ci_stage, pipeline: pipeline, project: pipeline.project) }
 
   it_behaves_like 'having unique enum values'
 
@@ -30,9 +30,9 @@ RSpec.describe Ci::Stage, :models do
 
   describe '.by_position' do
     it 'finds stages by position' do
-      a = create(:ci_stage_entity, position: 1)
-      b = create(:ci_stage_entity, position: 2)
-      c = create(:ci_stage_entity, position: 3)
+      a = create(:ci_stage, position: 1)
+      b = create(:ci_stage, position: 2)
+      c = create(:ci_stage, position: 3)
 
       expect(described_class.by_position(1)).to contain_exactly(a)
       expect(described_class.by_position(2)).to contain_exactly(b)
@@ -42,9 +42,9 @@ RSpec.describe Ci::Stage, :models do
 
   describe '.by_name' do
     it 'finds stages by name' do
-      a = create(:ci_stage_entity, name: 'a')
-      b = create(:ci_stage_entity, name: 'b')
-      c = create(:ci_stage_entity, name: 'c')
+      a = create(:ci_stage, name: 'a')
+      b = create(:ci_stage, name: 'b')
+      c = create(:ci_stage, name: 'c')
 
       expect(described_class.by_name('a')).to contain_exactly(a)
       expect(described_class.by_name('b')).to contain_exactly(b)
@@ -54,7 +54,7 @@ RSpec.describe Ci::Stage, :models do
 
   describe '#status' do
     context 'when stage is pending' do
-      let(:stage) { create(:ci_stage_entity, status: 'pending') }
+      let(:stage) { create(:ci_stage, status: 'pending') }
 
       it 'has a correct status value' do
         expect(stage.status).to eq 'pending'
@@ -62,7 +62,7 @@ RSpec.describe Ci::Stage, :models do
     end
 
     context 'when stage is success' do
-      let(:stage) { create(:ci_stage_entity, status: 'success') }
+      let(:stage) { create(:ci_stage, status: 'success') }
 
       it 'has a correct status value' do
         expect(stage.status).to eq 'success'
@@ -119,7 +119,7 @@ RSpec.describe Ci::Stage, :models do
     end
 
     context 'when stage has only created builds' do
-      let(:stage) { create(:ci_stage_entity, status: :created) }
+      let(:stage) { create(:ci_stage, status: :created) }
 
       before do
         create(:ci_build, :created, stage_id: stage.id)
@@ -206,7 +206,7 @@ RSpec.describe Ci::Stage, :models do
     using RSpec::Parameterized::TableSyntax
 
     let(:user) { create(:user) }
-    let(:stage) { create(:ci_stage_entity, status: :created) }
+    let(:stage) { create(:ci_stage, status: :created) }
 
     subject { stage.detailed_status(user) }
 
@@ -269,7 +269,7 @@ RSpec.describe Ci::Stage, :models do
   describe '#delay' do
     subject { stage.delay }
 
-    let(:stage) { create(:ci_stage_entity, status: :created) }
+    let(:stage) { create(:ci_stage, status: :created) }
 
     it 'updates stage status' do
       subject
@@ -361,12 +361,12 @@ RSpec.describe Ci::Stage, :models do
     end
   end
 
-  it_behaves_like 'manual playable stage', :ci_stage_entity
+  it_behaves_like 'manual playable stage', :ci_stage
 
   context 'loose foreign key on ci_stages.project_id' do
     it_behaves_like 'cleanup by a loose foreign key' do
       let!(:parent) { create(:project) }
-      let!(:model) { create(:ci_stage_entity, project: parent) }
+      let!(:model) { create(:ci_stage, project: parent) }
     end
   end
 end
diff --git a/spec/models/commit_status_spec.rb b/spec/models/commit_status_spec.rb
index 9f58d3c1a80..3cccc41a892 100644
--- a/spec/models/commit_status_spec.rb
+++ b/spec/models/commit_status_spec.rb
@@ -803,7 +803,7 @@ RSpec.describe CommitStatus do
   describe 'ensure stage assignment' do
     context 'when commit status has a stage_id assigned' do
       let!(:stage) do
-        create(:ci_stage_entity, project: project, pipeline: pipeline)
+        create(:ci_stage, project: project, pipeline: pipeline)
       end
 
       let(:commit_status) do
@@ -836,7 +836,7 @@ RSpec.describe CommitStatus do
 
     context 'when commit status does not have stage but it exists' do
       let!(:stage) do
-        create(:ci_stage_entity, project: project,
+        create(:ci_stage, project: project,
                                  pipeline: pipeline,
                                  name: 'test')
       end
diff --git a/spec/presenters/ci/stage_presenter_spec.rb b/spec/presenters/ci/stage_presenter_spec.rb
index 6b82785c2d4..368f03b0150 100644
--- a/spec/presenters/ci/stage_presenter_spec.rb
+++ b/spec/presenters/ci/stage_presenter_spec.rb
@@ -3,7 +3,7 @@
 require 'spec_helper'
 
 RSpec.describe Ci::StagePresenter do
-  let(:stage) { create(:ci_stage_entity) }
+  let(:stage) { create(:ci_stage) }
   let(:presenter) { described_class.new(stage) }
 
   let!(:build) { create(:ci_build, :tags, :artifacts, pipeline: stage.pipeline, stage: stage.name) }
diff --git a/spec/requests/api/graphql/ci/job_spec.rb b/spec/requests/api/graphql/ci/job_spec.rb
index 2fb90dcd92b..3721155c71b 100644
--- a/spec/requests/api/graphql/ci/job_spec.rb
+++ b/spec/requests/api/graphql/ci/job_spec.rb
@@ -13,8 +13,8 @@ RSpec.describe 'Query.project(fullPath).pipelines.job(id)' do
   let_it_be(:project) { create(:project, :repository, :public) }
   let_it_be(:pipeline) { create(:ci_pipeline, project: project) }
 
-  let_it_be(:prepare_stage) { create(:ci_stage_entity, pipeline: pipeline, project: project, name: 'prepare') }
-  let_it_be(:test_stage) { create(:ci_stage_entity, pipeline: pipeline, project: project, name: 'test') }
+  let_it_be(:prepare_stage) { create(:ci_stage, pipeline: pipeline, project: project, name: 'prepare') }
+  let_it_be(:test_stage) { create(:ci_stage, pipeline: pipeline, project: project, name: 'test') }
 
   let_it_be(:job_1) { create(:ci_build, pipeline: pipeline, stage: 'prepare', name: 'Job 1') }
   let_it_be(:job_2) { create(:ci_build, pipeline: pipeline, stage: 'test', name: 'Job 2') }
diff --git a/spec/requests/api/graphql/ci/jobs_spec.rb b/spec/requests/api/graphql/ci/jobs_spec.rb
index d1737fc22ae..8c4ab13fc35 100644
--- a/spec/requests/api/graphql/ci/jobs_spec.rb
+++ b/spec/requests/api/graphql/ci/jobs_spec.rb
@@ -14,7 +14,7 @@ RSpec.describe 'Query.project.pipeline' do
   describe '.stages.groups.jobs' do
     let(:pipeline) do
       pipeline = create(:ci_pipeline, project: project, user: user)
-      stage = create(:ci_stage_entity, project: project, pipeline: pipeline, name: 'first', position: 1)
+      stage = create(:ci_stage, project: project, pipeline: pipeline, name: 'first', position: 1)
       create(:ci_build, stage_id: stage.id, pipeline: pipeline, name: 'my test job', scheduling_type: :stage)
 
       pipeline
@@ -84,8 +84,8 @@ RSpec.describe 'Query.project.pipeline' do
 
     context 'when there is more than one stage and job needs' do
       before do
-        build_stage = create(:ci_stage_entity, position: 2, name: 'build', project: project, pipeline: pipeline)
-        test_stage = create(:ci_stage_entity, position: 3, name: 'test', project: project, pipeline: pipeline)
+        build_stage = create(:ci_stage, position: 2, name: 'build', project: project, pipeline: pipeline)
+        test_stage = create(:ci_stage, position: 3, name: 'test', project: project, pipeline: pipeline)
 
         create(:ci_build, pipeline: pipeline, name: 'docker 1 2', scheduling_type: :stage, stage: build_stage, stage_idx: build_stage.position)
         create(:ci_build, pipeline: pipeline, name: 'docker 2 2', stage: build_stage, stage_idx: build_stage.position, scheduling_type: :dag)
diff --git a/spec/requests/api/graphql/ci/pipelines_spec.rb b/spec/requests/api/graphql/ci/pipelines_spec.rb
index 741af676b6d..a968e5508cb 100644
--- a/spec/requests/api/graphql/ci/pipelines_spec.rb
+++ b/spec/requests/api/graphql/ci/pipelines_spec.rb
@@ -86,8 +86,8 @@ RSpec.describe 'Query.project(fullPath).pipelines' do
   describe '.stages' do
     let_it_be(:project) { create(:project, :repository) }
     let_it_be(:pipeline) { create(:ci_empty_pipeline, project: project) }
-    let_it_be(:stage) { create(:ci_stage_entity, pipeline: pipeline, project: project) }
-    let_it_be(:other_stage) { create(:ci_stage_entity, pipeline: pipeline, project: project, name: 'other') }
+    let_it_be(:stage) { create(:ci_stage, pipeline: pipeline, project: project) }
+    let_it_be(:other_stage) { create(:ci_stage, pipeline: pipeline, project: project, name: 'other') }
 
     let(:first_n) { var('Int') }
     let(:query_path) do
diff --git a/spec/requests/api/graphql/ci/stages_spec.rb b/spec/requests/api/graphql/ci/stages_spec.rb
index 50d2cf75097..1edd6e58486 100644
--- a/spec/requests/api/graphql/ci/stages_spec.rb
+++ b/spec/requests/api/graphql/ci/stages_spec.rb
@@ -36,7 +36,7 @@ RSpec.describe 'Query.project.pipeline.stages' do
   end
 
   before_all do
-    create(:ci_stage_entity, pipeline: pipeline, name: 'deploy')
+    create(:ci_stage, pipeline: pipeline, name: 'deploy')
     create_list(:ci_build, 2, pipeline: pipeline, stage: 'deploy')
   end
 
diff --git a/spec/requests/api/graphql/project/jobs_spec.rb b/spec/requests/api/graphql/project/jobs_spec.rb
index 1a823ede9ac..7d0eb203d60 100644
--- a/spec/requests/api/graphql/project/jobs_spec.rb
+++ b/spec/requests/api/graphql/project/jobs_spec.rb
@@ -31,8 +31,8 @@ RSpec.describe 'Query.project.jobs' do
   end
 
   it 'does not generate N+1 queries', :request_store, :use_sql_query_cache do
-    build_stage = create(:ci_stage_entity, position: 1, name: 'build', project: project, pipeline: pipeline)
-    test_stage = create(:ci_stage_entity, position: 2, name: 'test', project: project, pipeline: pipeline)
+    build_stage = create(:ci_stage, position: 1, name: 'build', project: project, pipeline: pipeline)
+    test_stage = create(:ci_stage, position: 2, name: 'test', project: project, pipeline: pipeline)
     create(:ci_build, pipeline: pipeline, stage_idx: build_stage.position, name: 'docker 1 2', stage: build_stage)
     create(:ci_build, pipeline: pipeline, stage_idx: build_stage.position, name: 'docker 2 2', stage: build_stage)
     create(:ci_build, pipeline: pipeline, stage_idx: test_stage.position, name: 'rspec 1 2', stage: test_stage)
diff --git a/spec/requests/api/graphql/project/pipeline_spec.rb b/spec/requests/api/graphql/project/pipeline_spec.rb
index ccf97918021..1151438b446 100644
--- a/spec/requests/api/graphql/project/pipeline_spec.rb
+++ b/spec/requests/api/graphql/project/pipeline_spec.rb
@@ -290,8 +290,8 @@ RSpec.describe 'getting pipeline information nested in a project' do
     end
 
     it 'does not generate N+1 queries', :request_store, :use_sql_query_cache do
-      build_stage = create(:ci_stage_entity, position: 1, name: 'build', project: project, pipeline: pipeline)
-      test_stage = create(:ci_stage_entity, position: 2, name: 'test', project: project, pipeline: pipeline)
+      build_stage = create(:ci_stage, position: 1, name: 'build', project: project, pipeline: pipeline)
+      test_stage = create(:ci_stage, position: 2, name: 'test', project: project, pipeline: pipeline)
       create(:ci_build, pipeline: pipeline, stage_idx: build_stage.position, name: 'docker 1 2', stage: build_stage)
       create(:ci_build, pipeline: pipeline, stage_idx: build_stage.position, name: 'docker 2 2', stage: build_stage)
       create(:ci_build, pipeline: pipeline, stage_idx: test_stage.position, name: 'rspec 1 2', stage: test_stage)
diff --git a/spec/serializers/ci/dag_job_group_entity_spec.rb b/spec/serializers/ci/dag_job_group_entity_spec.rb
index d05dd8b05b9..b654b21f583 100644
--- a/spec/serializers/ci/dag_job_group_entity_spec.rb
+++ b/spec/serializers/ci/dag_job_group_entity_spec.rb
@@ -5,7 +5,7 @@ require 'spec_helper'
 RSpec.describe Ci::DagJobGroupEntity do
   let_it_be(:request) { double(:request) }
   let_it_be(:pipeline) { create(:ci_pipeline) }
-  let_it_be(:stage) { create(:ci_stage_entity, pipeline: pipeline) }
+  let_it_be(:stage) { create(:ci_stage, pipeline: pipeline) }
 
   let(:group) { Ci::Group.new(pipeline.project, stage, name: 'test', jobs: jobs) }
   let(:entity) { described_class.new(group, request: request) }
diff --git a/spec/serializers/ci/dag_pipeline_entity_spec.rb b/spec/serializers/ci/dag_pipeline_entity_spec.rb
index 31a0dc5c048..548fd247743 100644
--- a/spec/serializers/ci/dag_pipeline_entity_spec.rb
+++ b/spec/serializers/ci/dag_pipeline_entity_spec.rb
@@ -43,9 +43,9 @@ RSpec.describe Ci::DagPipelineEntity do
     end
 
     context 'when pipeline has parallel jobs, DAG needs and GenericCommitStatus' do
-      let!(:stage_build)  { create(:ci_stage_entity, name: 'build',  position: 1, pipeline: pipeline) }
-      let!(:stage_test)   { create(:ci_stage_entity, name: 'test',   position: 2, pipeline: pipeline) }
-      let!(:stage_deploy) { create(:ci_stage_entity, name: 'deploy', position: 3, pipeline: pipeline) }
+      let!(:stage_build)  { create(:ci_stage, name: 'build',  position: 1, pipeline: pipeline) }
+      let!(:stage_test)   { create(:ci_stage, name: 'test',   position: 2, pipeline: pipeline) }
+      let!(:stage_deploy) { create(:ci_stage, name: 'deploy', position: 3, pipeline: pipeline) }
 
       let!(:job_build_1)   { create(:ci_build, name: 'build 1', stage: 'build', pipeline: pipeline) }
       let!(:job_build_2)   { create(:ci_build, name: 'build 2', stage: 'build', pipeline: pipeline) }
diff --git a/spec/serializers/ci/dag_stage_entity_spec.rb b/spec/serializers/ci/dag_stage_entity_spec.rb
index 554437ecef5..3530a5e2bae 100644
--- a/spec/serializers/ci/dag_stage_entity_spec.rb
+++ b/spec/serializers/ci/dag_stage_entity_spec.rb
@@ -6,7 +6,7 @@ RSpec.describe Ci::DagStageEntity do
   let_it_be(:pipeline) { create(:ci_pipeline) }
   let_it_be(:request) { double(:request) }
 
-  let(:stage) { create(:ci_stage_entity, pipeline: pipeline, name: 'test') }
+  let(:stage) { create(:ci_stage, pipeline: pipeline, name: 'test') }
   let(:entity) { described_class.new(stage, request: request) }
 
   let!(:job) { create(:ci_build, :success, pipeline: pipeline, stage_id: stage.id) }
diff --git a/spec/serializers/stage_entity_spec.rb b/spec/serializers/stage_entity_spec.rb
index ab74ce76cbd..95d3fd254d4 100644
--- a/spec/serializers/stage_entity_spec.rb
+++ b/spec/serializers/stage_entity_spec.rb
@@ -12,7 +12,7 @@ RSpec.describe StageEntity do
   end
 
   let(:stage) do
-    create(:ci_stage_entity, pipeline: pipeline, status: :success)
+    create(:ci_stage, pipeline: pipeline, status: :success)
   end
 
   before do
@@ -74,7 +74,7 @@ RSpec.describe StageEntity do
     end
 
     context 'with a skipped stage ' do
-      let(:stage) { create(:ci_stage_entity, status: 'skipped') }
+      let(:stage) { create(:ci_stage, status: 'skipped') }
 
       it 'contains play_all_manual' do
         expect(subject[:status][:action]).to be_present
@@ -82,7 +82,7 @@ RSpec.describe StageEntity do
     end
 
     context 'with a scheduled stage ' do
-      let(:stage) { create(:ci_stage_entity, status: 'scheduled') }
+      let(:stage) { create(:ci_stage, status: 'scheduled') }
 
       it 'contains play_all_manual' do
         expect(subject[:status][:action]).to be_present
@@ -90,7 +90,7 @@ RSpec.describe StageEntity do
     end
 
     context 'with a manual stage ' do
-      let(:stage) { create(:ci_stage_entity, status: 'manual') }
+      let(:stage) { create(:ci_stage, status: 'manual') }
 
       it 'contains play_all_manual' do
         expect(subject[:status][:action]).to be_present
diff --git a/spec/serializers/stage_serializer_spec.rb b/spec/serializers/stage_serializer_spec.rb
index 0b5e87dc95b..24e8057375b 100644
--- a/spec/serializers/stage_serializer_spec.rb
+++ b/spec/serializers/stage_serializer_spec.rb
@@ -5,7 +5,7 @@ require 'spec_helper'
 RSpec.describe StageSerializer do
   let(:project) { create(:project, :repository) }
   let(:user) { create(:user) }
-  let(:resource) { create(:ci_stage_entity) }
+  let(:resource) { create(:ci_stage) }
 
   let(:serializer) do
     described_class.new(current_user: user, project: project)
@@ -21,7 +21,7 @@ RSpec.describe StageSerializer do
     end
 
     context 'with an array of entities' do
-      let(:resource) { create_list(:ci_stage_entity, 2) }
+      let(:resource) { create_list(:ci_stage, 2) }
 
       it 'serializes the array of pipelines' do
         expect(subject).not_to be_empty
diff --git a/spec/services/ci/abort_pipelines_service_spec.rb b/spec/services/ci/abort_pipelines_service_spec.rb
index 9f9519d6829..e43faf0af51 100644
--- a/spec/services/ci/abort_pipelines_service_spec.rb
+++ b/spec/services/ci/abort_pipelines_service_spec.rb
@@ -12,13 +12,13 @@ RSpec.describe Ci::AbortPipelinesService do
 
   let_it_be(:cancelable_build, reload: true) { create(:ci_build, :running, pipeline: cancelable_pipeline) }
   let_it_be(:non_cancelable_build, reload: true) { create(:ci_build, :success, pipeline: cancelable_pipeline) }
-  let_it_be(:cancelable_stage, reload: true) { create(:ci_stage_entity, name: 'stageA', status: :running, pipeline: cancelable_pipeline, project: project) }
-  let_it_be(:non_cancelable_stage, reload: true) { create(:ci_stage_entity, name: 'stageB', status: :success, pipeline: cancelable_pipeline, project: project) }
+  let_it_be(:cancelable_stage, reload: true) { create(:ci_stage, name: 'stageA', status: :running, pipeline: cancelable_pipeline, project: project) }
+  let_it_be(:non_cancelable_stage, reload: true) { create(:ci_stage, name: 'stageB', status: :success, pipeline: cancelable_pipeline, project: project) }
 
   let_it_be(:manual_pipeline_cancelable_build, reload: true) { create(:ci_build, :created, pipeline: manual_pipeline) }
   let_it_be(:manual_pipeline_non_cancelable_build, reload: true) { create(:ci_build, :manual, pipeline: manual_pipeline) }
-  let_it_be(:manual_pipeline_cancelable_stage, reload: true) { create(:ci_stage_entity, name: 'stageA', status: :created, pipeline: manual_pipeline, project: project) }
-  let_it_be(:manual_pipeline_non_cancelable_stage, reload: true) { create(:ci_stage_entity, name: 'stageB', status: :success, pipeline: manual_pipeline, project: project) }
+  let_it_be(:manual_pipeline_cancelable_stage, reload: true) { create(:ci_stage, name: 'stageA', status: :created, pipeline: manual_pipeline, project: project) }
+  let_it_be(:manual_pipeline_non_cancelable_stage, reload: true) { create(:ci_stage, name: 'stageB', status: :success, pipeline: manual_pipeline, project: project) }
 
   describe '#execute' do
     def expect_correct_pipeline_cancellations
diff --git a/spec/services/ci/ensure_stage_service_spec.rb b/spec/services/ci/ensure_stage_service_spec.rb
index 3ede214cdd4..026814edda6 100644
--- a/spec/services/ci/ensure_stage_service_spec.rb
+++ b/spec/services/ci/ensure_stage_service_spec.rb
@@ -6,7 +6,7 @@ RSpec.describe Ci::EnsureStageService, '#execute' do
   let_it_be(:project) { create(:project) }
   let_it_be(:user) { create(:user) }
 
-  let(:stage) { create(:ci_stage_entity) }
+  let(:stage) { create(:ci_stage) }
   let(:job) { build(:ci_build) }
 
   let(:service) { described_class.new(project, user) }
diff --git a/spec/services/ci/play_manual_stage_service_spec.rb b/spec/services/ci/play_manual_stage_service_spec.rb
index 3e2a95ee975..b3ae92aa787 100644
--- a/spec/services/ci/play_manual_stage_service_spec.rb
+++ b/spec/services/ci/play_manual_stage_service_spec.rb
@@ -11,7 +11,7 @@ RSpec.describe Ci::PlayManualStageService, '#execute' do
   let(:stage_status) { 'manual' }
 
   let(:stage) do
-    create(:ci_stage_entity,
+    create(:ci_stage,
            pipeline: pipeline,
            project: project,
            name: 'test')
diff --git a/spec/services/ci/retry_job_service_spec.rb b/spec/services/ci/retry_job_service_spec.rb
index acc7a99637b..f042471bd1f 100644
--- a/spec/services/ci/retry_job_service_spec.rb
+++ b/spec/services/ci/retry_job_service_spec.rb
@@ -12,7 +12,7 @@ RSpec.describe Ci::RetryJobService do
   end
 
   let_it_be(:stage) do
-    create(:ci_stage_entity, project: project,
+    create(:ci_stage, project: project,
                              pipeline: pipeline,
                              name: 'test')
   end
@@ -154,7 +154,7 @@ RSpec.describe Ci::RetryJobService do
     end
 
     context 'when the pipeline has other jobs' do
-      let!(:stage2) { create(:ci_stage_entity, project: project, pipeline: pipeline, name: 'deploy') }
+      let!(:stage2) { create(:ci_stage, project: project, pipeline: pipeline, name: 'deploy') }
       let!(:build2) { create(:ci_build, pipeline: pipeline, stage_id: stage.id ) }
       let!(:deploy) { create(:ci_build, pipeline: pipeline, stage_id: stage2.id) }
       let!(:deploy_needs_build2) { create(:ci_build_need, build: deploy, name: build2.name) }
diff --git a/spec/workers/stage_update_worker_spec.rb b/spec/workers/stage_update_worker_spec.rb
index 75b324a9e0a..e50c7183153 100644
--- a/spec/workers/stage_update_worker_spec.rb
+++ b/spec/workers/stage_update_worker_spec.rb
@@ -5,7 +5,7 @@ require 'spec_helper'
 RSpec.describe StageUpdateWorker do
   describe '#perform' do
     context 'when stage exists' do
-      let(:stage) { create(:ci_stage_entity) }
+      let(:stage) { create(:ci_stage) }
 
       it 'updates stage status' do
         expect_any_instance_of(Ci::Stage).to receive(:set_status).with('skipped')