Merge branch 'avoid-n-plus-1-audit-events' into 'master'

Avoid N+1 queries on Audit Events controllers

Closes #212839

See merge request gitlab-org/gitlab!28399

app/models/audit_event.rb

@@ -30,12 +30,26 @@ class AuditEvent < ApplicationRecord
   end
 
   def author_name
-    self.user.name
+    lazy_author.name
   end
 
   def formatted_details
     details.merge(details.slice(:from, :to).transform_values(&:to_s))
   end
+
+  def lazy_author
+    BatchLoader.for(author_id).batch(default_value: default_author_value) do |author_ids, loader|
+      User.where(id: author_ids).find_each do |user|
+        loader.call(user.id, user)
+      end
+    end
+  end
+
+  private
+
+  def default_author_value
+    ::Gitlab::Audit::NullAuthor.for(author_id, details[:author_name])
+  end
 end
 
 AuditEvent.prepend_if_ee('EE::AuditEvent')

app/services/audit_event_service.rb

@@ -13,7 +13,7 @@ class AuditEventService
   #
   # @return [AuditEventService]
   def initialize(author, entity, details = {})
-    @author = author
+    @author = build_author(author)
     @entity = entity
     @details = details
   end
@@ -49,6 +49,14 @@ class AuditEventService
 
   private
 
+  def build_author(author)
+    if author.is_a?(User)
+      author
+    else
+      Gitlab::Audit::UnauthenticatedAuthor.new(name: author)
+    end
+  end
+
   def base_payload
     {
       author_id: @author.id,

ee/app/controllers/admin/audit_logs_controller.rb

@@ -10,7 +10,7 @@ class Admin::AuditLogsController < Admin::ApplicationController
   PER_PAGE = 25
 
   def index
-    @events = audit_log_events.page(params[:page]).per(PER_PAGE).without_count
+    @events = audit_log_events
 
     @entity = case audit_logs_params[:entity_type]
               when 'User'
@@ -27,7 +27,10 @@ class Admin::AuditLogsController < Admin::ApplicationController
   private
 
   def audit_log_events
-    AuditLogFinder.new(audit_logs_params).execute
+    events = AuditLogFinder.new(audit_logs_params).execute
+    events = events.page(params[:page]).per(PER_PAGE).without_count
+
+    Gitlab::Audit::Events::Preloader.preload!(events)
   end
 
   def check_license_admin_audit_log_available!

ee/app/controllers/groups/audit_events_controller.rb

@@ -11,7 +11,9 @@ class Groups::AuditEventsController < Groups::ApplicationController
   layout 'group_settings'
 
   def index
-    @events = AuditLogFinder.new(audit_logs_params).execute.page(params[:page])
+    events = AuditLogFinder.new(audit_logs_params).execute.page(params[:page])
+
+    @events = Gitlab::Audit::Events::Preloader.preload!(events)
   end
 
   private

ee/app/controllers/projects/audit_events_controller.rb

@@ -12,7 +12,9 @@ class Projects::AuditEventsController < Projects::ApplicationController
   layout 'project_settings'
 
   def index
-    @events = AuditLogFinder.new(audit_logs_params).execute.page(params[:page])
+    events = AuditLogFinder.new(audit_logs_params).execute.page(params[:page])
+
+    @events = Gitlab::Audit::Events::Preloader.preload!(events)
   end
 
   private

ee/app/models/ee/audit_event.rb

@@ -5,36 +5,22 @@ module EE
     extend ActiveSupport::Concern
     extend ::Gitlab::Utils::Override
 
-    # While tracking events that could take place even when
-    # a user is not logged in, (eg: downloading repo of a public project),
-    # we set the author_id of such events as -1
-    UNAUTH_USER_AUTHOR_ID = -1
-    # Events that are authored by unathenticated users, should be
-    # shown as authored by `An unauthenticated user` in the UI.
-    UNAUTH_USER_AUTHOR_NAME = 'An unauthenticated user'.freeze
-
-    override :author_name
-    def author_name
-      if (author_name = details[:author_name].presence || user&.name)
-        author_name
-      elsif authored_by_unauth_user?
-        UNAUTH_USER_AUTHOR_NAME
-      end
-    end
-
     def entity
-      return unless entity_type && entity_id
-
-      # Avoiding exception if the record doesn't exist
-      @entity ||= entity_type.constantize.find_by_id(entity_id) # rubocop:disable Gitlab/ModuleWithInstanceVariables
+      lazy_entity
     end
 
     def present
       AuditEventPresenter.new(self)
     end
 
-    def authored_by_unauth_user?
-      author_id == UNAUTH_USER_AUTHOR_ID
+    def lazy_entity
+      BatchLoader.for(entity_id)
+        .batch(
+          key: entity_type, default_value: ::Gitlab::Audit::NullEntity.new
+        ) do |ids, loader, args|
+          model = Object.const_get(args[:key], false)
+          model.where(id: ids).find_each { |record| loader.call(record.id, record) }
+        end
     end
   end
 end

ee/app/presenters/audit_event_presenter.rb

@@ -4,13 +4,11 @@ class AuditEventPresenter < Gitlab::View::Presenter::Simple
   presents :audit_event
 
   def author_name
-    user = audit_event.user
+    author = audit_event.lazy_author
 
-    if user
-      link_to(user.name, user_path(user))
-    else
-      audit_event.author_name
-    end
+    return author.name if author.is_a?(Gitlab::Audit::NullAuthor)
+
+    link_to(author.name, user_path(author))
   end
 
   def target
@@ -26,9 +24,9 @@ class AuditEventPresenter < Gitlab::View::Presenter::Simple
   end
 
   def object
-    entity = audit_event.entity
+    entity = audit_event.lazy_entity
 
-    return unless entity
+    return if entity.is_a?(Gitlab::Audit::NullEntity)
 
     link_to(details[:entity_path] || entity.name, entity).html_safe
   end

ee/app/services/ee/audit_event_service.rb

@@ -89,7 +89,7 @@ module EE
 
       @details = {
         failed_login: auth.upcase,
-        author_name: @author,
+        author_name: @author.name,
         target_details: @author,
         ip_address: ip
       }
@@ -141,7 +141,7 @@ module EE
       @details[:entity_path] = @entity&.full_path if admin_audit_log_enabled?
 
       SecurityEvent.create(
-        author_id: @author.respond_to?(:id) ? @author.id : AuditEvent::UNAUTH_USER_AUTHOR_ID,
+        author_id: @author.id,
         entity_id: @entity.respond_to?(:id) ? @entity.id : -1,
         entity_type: 'User',
         details: @details
@@ -213,7 +213,7 @@ module EE
     override :base_payload
     def base_payload
       {
-        author_id: @author.respond_to?(:id) ? @author.id : AuditEvent::UNAUTH_USER_AUTHOR_ID,
+        author_id: @author.id,
         # `@author.respond_to?(:id)` is to support cases where we need to log events
         # that could take place even when a user is unathenticated, Eg: downloading a public repo.
         # For such events, it is not mandatory that an author is always present.
@@ -260,7 +260,7 @@ module EE
     end
 
     def ip_address
-      @author&.current_sign_in_ip || @details[:ip_address]
+      @author.current_sign_in_ip || @details[:ip_address]
     end
 
     def add_security_event_admin_details!

ee/changelogs/unreleased/fix-audit-events-controllers-n-plus-1.yml

+---
+title: Fix N+1 queries in Audit Events controllers
+merge_request: 28399
+author:
+type: performance

ee/lib/gitlab/audit/deleted_author.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Audit
+    class DeletedAuthor < Gitlab::Audit::NullAuthor
+    end
+  end
+end

ee/lib/gitlab/audit/events/preloader.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Audit
+    module Events
+      class Preloader
+        def self.preload!(audit_events)
+          audit_events.tap do |audit_events|
+            audit_events.each do |audit_event|
+              audit_event.lazy_author
+              audit_event.lazy_entity
+            end
+          end
+        end
+      end
+    end
+  end
+end

ee/lib/gitlab/audit/null_author.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Audit
+    class NullAuthor
+      attr_reader :id, :name
+
+      # Creates an Author
+      #
+      # While tracking events that could take place even when
+      # a user is not logged in, (eg: downloading repo of a public project),
+      # we set the author_id of such events as -1
+      #
+      # @param [Integer] id
+      # @param [String] name
+      #
+      # @return [Gitlab::Audit::UnauthenticatedAuthor, Gitlab::Audit::DeletedAuthor]
+      def self.for(id, name)
+        if id == -1
+          Gitlab::Audit::UnauthenticatedAuthor.new(name: name)
+        else
+          Gitlab::Audit::DeletedAuthor.new(id: id, name: name)
+        end
+      end
+
+      def initialize(id:, name:)
+        @id = id
+        @name = name
+      end
+
+      def current_sign_in_ip
+        nil
+      end
+    end
+  end
+end

ee/lib/gitlab/audit/null_entity.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Audit
+    class NullEntity
+    end
+  end
+end

ee/lib/gitlab/audit/unauthenticated_author.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Audit
+    class UnauthenticatedAuthor < Gitlab::Audit::NullAuthor
+      def initialize(name: nil)
+        super(id: -1, name: name)
+      end
+
+      # Events that are authored by unathenticated users, should be
+      # shown as authored by `An unauthenticated user` in the UI.
+      def name
+        @name || 'An unauthenticated user'
+      end
+    end
+  end
+end

ee/spec/factories/audit_events.rb

@@ -4,33 +4,38 @@ FactoryBot.define do
   factory :audit_event, class: 'SecurityEvent', aliases: [:user_audit_event] do
     user
 
+    transient { target_user { create(:user) } }
+
     entity_type { 'User' }
-    entity_id   { user.id }
+    entity_id   { target_user.id }
     details do
       {
         change: 'email address',
         from: 'admin@gitlab.com',
         to: 'maintainer@gitlab.com',
         author_name: user.name,
-        target_id: user.id,
+        target_id: target_user.id,
         target_type: 'User',
-        target_details: user.name,
+        target_details: target_user.name,
         ip_address: '127.0.0.1',
-        entity_path: user.username
+        entity_path: target_user.username
       }
     end
 
     trait :project_event do
+      transient { target_project { create(:project) } }
+
       entity_type { 'Project' }
-      entity_id   { create(:project).id }
+      entity_id   { target_project.id }
       details do
         {
-          add: 'user_access',
-          as: 'Developer',
+          change: 'packges_enabled',
+          from: true,
+          to: false,
           author_name: user.name,
-          target_id: user.id,
-          target_type: 'User',
-          target_details: user.name,
+          target_id: target_project.id,
+          target_type: 'Project',
+          target_details: target_project.name,
           ip_address: '127.0.0.1',
           entity_path: 'gitlab.org/gitlab-ce'
         }
@@ -38,17 +43,19 @@ FactoryBot.define do
     end
 
     trait :group_event do
+      transient { target_group { create(:group) } }
+
       entity_type { 'Group' }
-      entity_id   { create(:group).id }
+      entity_id   { target_group.id }
       details do
         {
           change: 'project_creation_level',
           from: nil,
           to: 'Developers + Maintainers',
-          author_name: 'Administrator',
-          target_id: 1,
+          author_name: user.name,
+          target_id: target_group.id,
           target_type: 'Group',
-          target_details: "gitlab-org",
+          target_details: target_group.name,
           ip_address: '127.0.0.1',
           entity_path: "gitlab-org"
         }

ee/spec/lib/gitlab/audit/events/preloader_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Audit::Events::Preloader do
+  describe '.preload!' do
+    let_it_be(:audit_events) { create_list(:audit_event, 2) }
+
+    let(:audit_events_relation) { AuditEvent.where(id: audit_events.map(&:id)) }
+
+    subject { described_class.preload!(audit_events_relation) }
+
+    it 'returns an ActiveRecord::Relation' do
+      expect(subject).to be_an(ActiveRecord::Relation)
+    end
+
+    it 'preloads associated records' do
+      log = ActiveRecord::QueryRecorder.new do
+        subject.map do |event|
+          [event.author_name, event.lazy_entity.name]
+        end
+      end
+
+      # Expected queries when requesting for AuditEvent with associated records
+      #
+      # 1. On the audit_events table
+      #    SELECT "audit_events".* FROM "audit_events"
+      # 2. On the users table for author_name
+      #    SELECT "users".* FROM "users" WHERE "users"."id" IN (1, 3)
+      # 3. On the users table for entity name
+      #    SELECT "users".* FROM "users" WHERE "users"."id" IN (2, 4)
+      #
+      expect(log.count).to eq(3)
+    end
+  end
+end

ee/spec/lib/gitlab/audit/null_author_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Audit::NullAuthor do
+  subject { described_class }
+
+  describe '.for' do
+    it 'returns an DeletedAuthor' do
+      expect(subject.for(666, 'Old Hat')).to be_a(Gitlab::Audit::DeletedAuthor)
+    end
+
+    it 'returns an UnauthenticatedAuthor when id equals -1', :aggregate_failures do
+      expect(subject.for(-1, 'Frank')).to be_a(Gitlab::Audit::UnauthenticatedAuthor)
+      expect(subject.for(-1, 'Frank')).to have_attributes(id: -1, name: 'Frank')
+    end
+  end
+
+  describe '#current_sign_in_ip' do
+    it { expect(subject.new(id: 888, name: 'Guest').current_sign_in_ip).to be_nil }
+  end
+end

ee/spec/lib/gitlab/audit/unauthenticated_author_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Audit::UnauthenticatedAuthor do
+  describe '#initialize' do
+    it 'sets correct attributes' do
+      expect(described_class.new(name: 'Peppa Pig'))
+        .to have_attributes(id: -1, name: 'Peppa Pig')
+    end
+
+    it 'sets default name when it is not provided' do
+      expect(described_class.new)
+        .to have_attributes(id: -1, name: 'An unauthenticated user')
+    end
+  end
+end

ee/spec/models/audit_event_spec.rb

@@ -91,8 +91,8 @@ RSpec.describe AuditEvent, type: :model do
     context 'when entity does not exist' do
       subject(:event) { described_class.new(entity_id: non_existing_record_id, entity_type: 'User') }
 
-      it 'returns nil' do
-        expect(event.entity).to be_blank
+      it 'returns a NullEntity' do
+        expect(event.entity).to be_a(Gitlab::Audit::NullEntity)
       end
     end
   end

ee/spec/requests/admin/audit_events_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe 'view audit events' do
+  describe 'GET /audit_events' do
+    let_it_be(:admin) { create(:admin) }
+    let_it_be(:audit_event) { create(:user_audit_event) }
+
+    before do
+      stub_licensed_features(admin_audit_log: true)
+
+      login_as(admin)
+    end
+
+    it 'returns 200 response' do
+      send_request
+
+      expect(response).to have_gitlab_http_status(:ok)
+    end
+
+    it 'avoids N+1 DB queries', :request_store do
+      # warm up cache so these initial queries would not leak in our QueryRecorder
+      send_request
+
+      control = ActiveRecord::QueryRecorder.new(skip_cached: false) { send_request }
+
+      create_list(:user_audit_event, 2)
+
+      expect do
+        send_request
+      end.not_to exceed_all_query_limit(control)
+    end
+
+    def send_request
+      get admin_audit_logs_path
+    end
+  end
+end

ee/spec/requests/api/audit_events_spec.rb

@@ -107,7 +107,7 @@ describe API::AuditEvents do
 
             expect(response["id"]).to eq(user_audit_event.id)
             expect(response["author_id"]).to eq(user_audit_event.user.id)
-            expect(response["entity_id"]).to eq(user_audit_event.user.id)
+            expect(response["entity_id"]).to eq(user_audit_event.entity_id)
             expect(response["entity_type"]).to eq('User')
             expect(Time.parse(response["created_at"])).to be_like_time(user_audit_event.created_at)
             expect(details).to eq user_audit_event.formatted_details.with_indifferent_access
@@ -157,7 +157,7 @@ describe API::AuditEvents do
 
               expect(json_response["id"]).to eq(user_audit_event.id)
               expect(json_response["author_id"]).to eq(user_audit_event.user.id)
-              expect(json_response["entity_id"]).to eq(user_audit_event.user.id)
+              expect(json_response["entity_id"]).to eq(user_audit_event.entity_id)
               expect(json_response["entity_type"]).to eq('User')
               expect(Time.parse(json_response["created_at"])).to be_like_time(user_audit_event.created_at)
               expect(details).to eq user_audit_event.formatted_details.with_indifferent_access

ee/spec/requests/api/repositories_spec.rb

@@ -6,18 +6,28 @@ describe API::Repositories do
   let(:project) { create(:project, :repository) }
 
   describe "GET /projects/:id/repository/archive(.:format)?:sha" do
-    shared_examples 'logs the audit event' do
+    shared_examples 'an auditable and successful request' do
       let(:route) { "/projects/#{project.id}/repository/archive" }
 
+      before do
+        stub_licensed_features(admin_audit_log: true)
+      end
+
       it 'logs the audit event' do
         expect do
           get api(route, current_user)
         end.to change { SecurityEvent.count }.by(1)
       end
+
+      it 'sends the archive' do
+        get api(route, current_user)
+
+        expect(response).to have_gitlab_http_status(:ok)
+      end
     end
 
     context 'when unauthenticated', 'and project is public' do
-      it_behaves_like 'logs the audit event' do
+      it_behaves_like 'an auditable and successful request' do
         let(:project) { create(:project, :public, :repository) }
         let(:current_user) { nil }
       end
@@ -28,7 +38,7 @@ describe API::Repositories do
         project.add_developer(user)
       end
 
-      it_behaves_like 'logs the audit event' do
+      it_behaves_like 'an auditable and successful request' do
         let(:user) { create(:user) }
         let(:current_user) { user }
       end

ee/spec/requests/groups/audit_events_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe 'view audit events' do
+  describe 'GET /groups/:group/-/audit_events' do
+    let_it_be(:group) { create(:group) }
+    let_it_be(:user) { create(:user) }
+    let_it_be(:audit_event) { create(:group_audit_event, entity_id: group.id) }
+
+    before_all do
+      group.add_owner(user)
+    end
+
+    before do
+      stub_licensed_features(audit_events: true)
+
+      login_as(user)
+    end
+
+    it 'returns 200 response' do
+      send_request
+
+      expect(response).to have_gitlab_http_status(:ok)
+    end
+
+    it 'avoids N+1 DB queries', :request_store do
+      # warm up cache so these initial queries would not leak in our QueryRecorder
+      send_request
+
+      control = ActiveRecord::QueryRecorder.new(skip_cached: false) { send_request }
+
+      create_list(:group_audit_event, 2, entity_id: group.id)
+
+      expect do
+        send_request
+      end.not_to exceed_all_query_limit(control)
+    end
+
+    def send_request
+      get group_audit_events_path(group)
+    end
+  end
+end

ee/spec/requests/projects/audit_events_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe 'view audit events' do
+  describe 'GET /:namespace/:project/-/audit_events' do
+    let_it_be(:project) { create(:project, :repository) }
+    let_it_be(:user) { project.owner }
+    let_it_be(:audit_event) { create(:project_audit_event, entity_id: project.id) }
+
+    before do
+      stub_licensed_features(audit_events: true)
+
+      login_as(user)
+    end
+
+    it 'returns 200 response' do
+      send_request
+
+      expect(response).to have_gitlab_http_status(:ok)
+    end
+
+    it 'avoids N+1 DB queries', :request_store do
+      control = ActiveRecord::QueryRecorder.new(skip_cached: false) { send_request }
+
+      create_list(:project_audit_event, 2, entity_id: project.id)
+
+      expect do
+        send_request
+      end.not_to exceed_all_query_limit(control)
+    end
+
+    def send_request
+      get namespace_project_audit_events_path(
+        namespace_id: project.namespace,
+        project_id: project
+      )
+    end
+  end
+end

ee/spec/services/audit_event_service_spec.rb

@@ -4,9 +4,11 @@ require 'spec_helper'
 
 describe AuditEventService do
   let(:project) { create(:project) }
-  let(:user) { create(:user) }
+  let(:user) { create(:user, current_sign_in_ip: '192.168.68.104') }
   let(:project_member) { create(:project_member, user: user, expires_at: 1.day.from_now) }
-  let(:service) { described_class.new(user, project, { action: :destroy }) }
+
+  let(:details) { { action: :destroy } }
+  let(:service) { described_class.new(user, project, details) }
 
   describe '#for_member' do
     it 'generates event' do
@@ -21,11 +23,6 @@ describe AuditEventService do
       expect(event[:details][:target_details]).to eq('Deleted User')
     end
 
-    it 'has the IP address' do
-      event = service.for_member(project_member).security_event
-      expect(event[:details][:ip_address]).to eq(user.current_sign_in_ip)
-    end
-
     context 'user access expiry' do
       let(:service) { described_class.new(nil, project, { action: :expired }) }
 
@@ -63,21 +60,46 @@ describe AuditEventService do
 
       it 'has the entity full path' do
         event = service.for_member(project_member).security_event
+
         expect(event[:details][:entity_path]).to eq(project.full_path)
       end
+
+      it 'has the IP address' do
+        event = service.for_member(project_member).security_event
+
+        expect(event[:details][:ip_address]).to eq(user.current_sign_in_ip)
+      end
+    end
+
+    context 'admin audit log unlicensed' do
+      before do
+        stub_licensed_features(admin_audit_log: false)
+      end
+
+      it 'does not have the entity full path' do
+        event = service.for_member(project_member).security_event
+
+        expect(event[:details]).not_to have_key(:entity_path)
+      end
+
+      it 'does not have the ip_address' do
+        event = service.for_member(project_member).security_event
+
+        expect(event[:details]).not_to have_key(:ip_address)
+      end
     end
   end
 
   describe '#security_event' do
     context 'unlicensed' do
       before do
-        stub_licensed_features(audit_events: false)
+        disable_license_audit_features
       end
 
       it 'does not create an event' do
         expect(SecurityEvent).not_to receive(:create)
 
-        service.security_event
+        expect { service.security_event }.not_to change(SecurityEvent, :count)
       end
     end
 
@@ -96,6 +118,31 @@ describe AuditEventService do
         end
       end
     end
+
+    context 'admin audit log licensed' do
+      before do
+        stub_licensed_features(admin_audit_log: true)
+      end
+
+      context 'for an unauthenticated user' do
+        let(:details) { { ip_address: '10.11.12.13' } }
+        let(:user) { Gitlab::Audit::UnauthenticatedAuthor.new }
+
+        it 'defaults to the IP address in the details hash' do
+          event = service.security_event
+
+          expect(event[:details][:ip_address]).to eq('10.11.12.13')
+        end
+      end
+
+      context 'for an authenticated user' do
+        it 'has the user IP address' do
+          event = service.security_event
+
+          expect(event[:details][:ip_address]).to eq(user.current_sign_in_ip)
+        end
+      end
+    end
   end
 
   describe '#enabled?' do
@@ -110,9 +157,11 @@ describe AuditEventService do
 
     with_them do
       before do
-        stub_licensed_features(admin_audit_log: admin_audit_log,
-                               audit_events: audit_events,
-                               extended_audit_events: extended_audit_events)
+        stub_licensed_features(
+          admin_audit_log: admin_audit_log,
+          audit_events: audit_events,
+          extended_audit_events: extended_audit_events
+        )
       end
 
       it 'returns the correct result when feature is available' do
@@ -121,7 +170,7 @@ describe AuditEventService do
     end
   end
 
-  describe '#entity_audit_events_enabled??' do
+  describe '#entity_audit_events_enabled?' do
     context 'entity is a project' do
       let(:service) { described_class.new(user, project, { action: :destroy }) }
 
@@ -145,7 +194,7 @@ describe AuditEventService do
       it 'returns false when group is unlicensed' do
         stub_licensed_features(audit_events: false)
 
-        expect(service.entity_audit_events_enabled?).to be_falsy
+        expect(service.entity_audit_events_enabled?).to be_falsey
       end
 
       it 'returns true when group is licensed' do
@@ -200,23 +249,37 @@ describe AuditEventService do
       expect(event.details[:author_name]).to eq(author_name)
     end
 
-    it 'has the right IP address' do
-      allow(service).to receive(:admin_audit_log_enabled?).and_return(true)
-
-      expect(event.details[:ip_address]).to eq(ip_address)
-    end
-
     it 'has the right auth method for OAUTH' do
       oauth_service = described_class.new(author_name, nil, ip_address: ip_address, with: 'ldap')
       event = oauth_service.for_failed_login.unauth_security_event
 
       expect(event.details[:failed_login]).to eq('LDAP')
     end
+
+    context 'admin audit log licensed' do
+      before do
+        stub_licensed_features(extended_audit_events: true, admin_audit_log: true)
+      end
+
+      it 'has the right IP address' do
+        expect(event.details[:ip_address]).to eq(ip_address)
+      end
+    end
+
+    context 'admin audit log unlicensed' do
+      before do
+        stub_licensed_features(extended_audit_events: true, admin_audit_log: false)
+      end
+
+      it 'does not have the ip_address' do
+        expect(event.details).not_to have_key(:ip_address)
+      end
+    end
   end
 
   describe '#for_user' do
     let(:author_name) { 'Administrator' }
-    let(:current_user) { instance_spy(User, name: author_name) }
+    let(:current_user) { User.new(name: author_name) }
     let(:target_user_full_path) { 'ejohn' }
     let(:user) { instance_spy(User, full_path: target_user_full_path) }
     let(:custom_message) { 'Some strange event has occurred' }
@@ -291,18 +354,15 @@ describe AuditEventService do
   end
 
   describe 'license' do
-    let(:project) { create(:project) }
-    let(:user) { create(:user) }
-    let!(:service) { described_class.new(user, project, action: :create) }
     let(:event) { service.for_project.security_event }
 
     before do
-      disable_license_audit_features(service)
+      disable_license_audit_features
     end
 
     describe 'has the audit_admin feature' do
       before do
-        allow(service).to receive(:admin_audit_log_enabled?).and_return(true)
+        stub_licensed_features(admin_audit_log: true)
       end
 
       it 'logs an audit event' do
@@ -312,48 +372,60 @@ describe AuditEventService do
       it 'has the entity_path' do
         expect(event.details[:entity_path]).to eq(project.full_path)
       end
+
+      it 'has the user IP address' do
+        expect(event.details[:ip_address]).to eq(user.current_sign_in_ip)
+      end
     end
 
     describe 'has the extended_audit_events feature' do
       before do
-        allow(service).to receive(:entity_audit_events_enabled?).and_return(true)
+        stub_licensed_features(extended_audit_events: true)
       end
 
       it 'logs an audit event' do
         expect { event }.to change(AuditEvent, :count).by(1)
       end
 
-      it 'has not the entity_path' do
-        expect(event.details[:entity_path]).not_to eq(project.full_path)
+      it 'does not have the entity_path' do
+        expect(event.details).not_to have_key(:entity_path)
+      end
+
+      it 'does not have the ip_address' do
+        expect(event.details).not_to have_key(:ip_address)
       end
     end
 
     describe 'entity has the audit_events feature' do
       before do
-        allow(service).to receive(:audit_events_enabled?).and_return(true)
+        stub_licensed_features(audit_events: true)
       end
 
       it 'logs an audit event' do
         expect { event }.to change(AuditEvent, :count).by(1)
       end
 
-      it 'has not the entity_path' do
-        expect(event.details[:entity_path]).not_to eq(project.full_path)
+      it 'does not have the entity_path' do
+        expect(event.details).not_to have_key(:entity_path)
+      end
+
+      it 'does not have the ip_address' do
+        expect(event.details).not_to have_key(:ip_address)
       end
     end
 
-    describe 'has not any audit event feature' do
+    describe 'does not have any audit event feature' do
       it 'does not log the audit event' do
         expect { event }.not_to change(AuditEvent, :count)
       end
     end
+  end
 
-    def disable_license_audit_features(service)
-      [:entity_audit_events_enabled?,
-       :admin_audit_log_enabled?,
-       :audit_events_enabled?].each do |f|
-        allow(service).to receive(f).and_return(false)
-      end
-    end
+  def disable_license_audit_features
+    stub_licensed_features(
+      admin_audit_log: false,
+      audit_events: false,
+      extended_audit_events: false
+    )
   end
 end