Extend cop for verifying injecting of EE modules

This extends the InjectEnterpriseEditionModule RuboCop cop so that it
verifies the following:

1. The line number the injection occurs on (as before).
2. The method used (e.g. prepend instead of prepend_if_ee).
3. The argument type passed when using the new module injection methods.

2 files changed, 143 insertions, 49 deletions

diff --git a/rubocop/cop/inject_enterprise_edition_module.rb b/rubocop/cop/inject_enterprise_edition_module.rb
index 1d37b1bd12d..e0e1b2d6c7d 100644
--- a/rubocop/cop/inject_enterprise_edition_module.rb
+++ b/rubocop/cop/inject_enterprise_edition_module.rb
@@ -6,23 +6,39 @@ module RuboCop
     # the last line of a file. Injecting a module in the middle of a file will
     # cause merge conflicts, while placing it on the last line will not.
     class InjectEnterpriseEditionModule < RuboCop::Cop::Cop
-      MSG = 'Injecting EE modules must be done on the last line of this file' \
-        ', outside of any class or module definitions'
+      INVALID_LINE = 'Injecting EE modules must be done on the last line of this file' \
+          ', outside of any class or module definitions'
 
-      METHODS = Set.new(%i[include extend prepend]).freeze
+      DISALLOWED_METHOD =
+        'EE modules must be injected using `include_if_ee`, `extend_if_ee`, or `prepend_if_ee`'
+
+      INVALID_ARGUMENT = 'EE modules to inject must be specified as a String'
+
+      CHECK_LINE_METHODS =
+        Set.new(%i[include_if_ee extend_if_ee prepend_if_ee]).freeze
+
+      DISALLOW_METHODS = Set.new(%i[include extend prepend]).freeze
 
       def ee_const?(node)
         line = node.location.expression.source_line
 
         # We use `match?` here instead of RuboCop's AST matching, as this makes
         # it far easier to handle nested constants such as `EE::Foo::Bar::Baz`.
-        line.match?(/(\s|\()(::)?EE::/)
+        line.match?(/(\s|\()('|")?(::)?EE::/)
       end
 
       def on_send(node)
-        return unless METHODS.include?(node.children[1])
-        return unless ee_const?(node.children[2])
+        return unless check_method?(node)
 
+        if DISALLOW_METHODS.include?(node.children[1])
+          add_offense(node, message: DISALLOWED_METHOD)
+        else
+          verify_line_number(node)
+          verify_argument_type(node)
+        end
+      end
+
+      def verify_line_number(node)
         line = node.location.line
         buffer = node.location.expression.source_buffer
         last_line = buffer.last_line
@@ -32,7 +48,25 @@ module RuboCop
         # the expression is the last line _of code_.
         last_line -= 1 if buffer.source.end_with?("\n")
 
-        add_offense(node) if line < last_line
+        add_offense(node, message: INVALID_LINE) if line < last_line
+      end
+
+      def verify_argument_type(node)
+        argument = node.children[2]
+
+        return if argument.str_type?
+
+        add_offense(argument, message: INVALID_ARGUMENT)
+      end
+
+      def check_method?(node)
+        name = node.children[1]
+
+        if CHECK_LINE_METHODS.include?(name) || DISALLOW_METHODS.include?(name)
+          ee_const?(node.children[2])
+        else
+          false
+        end
       end
 
       # Automatically correcting these offenses is not always possible, as
diff --git a/spec/rubocop/cop/inject_enterprise_edition_module_spec.rb b/spec/rubocop/cop/inject_enterprise_edition_module_spec.rb
index 0ff777388e5..27df42c0aee 100644
--- a/spec/rubocop/cop/inject_enterprise_edition_module_spec.rb
+++ b/spec/rubocop/cop/inject_enterprise_edition_module_spec.rb
@@ -10,91 +10,91 @@ describe RuboCop::Cop::InjectEnterpriseEditionModule do
 
   subject(:cop) { described_class.new }
 
-  it 'flags the use of `prepend EE` in the middle of a file' do
+  it 'flags the use of `prepend_if_ee EE` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      prepend EE::Foo
-      ^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      prepend_if_ee 'EE::Foo'
+      ^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
 
-  it 'does not flag the use of `prepend EEFoo` in the middle of a file' do
+  it 'does not flag the use of `prepend_if_ee EEFoo` in the middle of a file' do
     expect_no_offenses(<<~SOURCE)
     class Foo
-      prepend EEFoo
+      prepend_if_ee 'EEFoo'
     end
     SOURCE
   end
 
-  it 'flags the use of `prepend EE::Foo::Bar` in the middle of a file' do
+  it 'flags the use of `prepend_if_ee EE::Foo::Bar` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      prepend EE::Foo::Bar
-      ^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      prepend_if_ee 'EE::Foo::Bar'
+      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
 
-  it 'flags the use of `prepend(EE::Foo::Bar)` in the middle of a file' do
+  it 'flags the use of `prepend_if_ee(EE::Foo::Bar)` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      prepend(EE::Foo::Bar)
-      ^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      prepend_if_ee('EE::Foo::Bar')
+      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
 
-  it 'flags the use of `prepend EE::Foo::Bar::Baz` in the middle of a file' do
+  it 'flags the use of `prepend_if_ee EE::Foo::Bar::Baz` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      prepend EE::Foo::Bar::Baz
-      ^^^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      prepend_if_ee 'EE::Foo::Bar::Baz'
+      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
 
-  it 'flags the use of `prepend ::EE` in the middle of a file' do
+  it 'flags the use of `prepend_if_ee ::EE` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      prepend ::EE::Foo
-      ^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      prepend_if_ee '::EE::Foo'
+      ^^^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
 
-  it 'flags the use of `include EE` in the middle of a file' do
+  it 'flags the use of `include_if_ee EE` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      include EE::Foo
-      ^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      include_if_ee 'EE::Foo'
+      ^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
 
-  it 'flags the use of `include ::EE` in the middle of a file' do
+  it 'flags the use of `include_if_ee ::EE` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      include ::EE::Foo
-      ^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      include_if_ee '::EE::Foo'
+      ^^^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
 
-  it 'flags the use of `extend EE` in the middle of a file' do
+  it 'flags the use of `extend_if_ee EE` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      extend EE::Foo
-      ^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      extend_if_ee 'EE::Foo'
+      ^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
 
-  it 'flags the use of `extend ::EE` in the middle of a file' do
+  it 'flags the use of `extend_if_ee ::EE` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      extend ::EE::Foo
-      ^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      extend_if_ee '::EE::Foo'
+      ^^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
@@ -102,7 +102,7 @@ describe RuboCop::Cop::InjectEnterpriseEditionModule do
   it 'does not flag prepending of regular modules' do
     expect_no_offenses(<<~SOURCE)
     class Foo
-      prepend Foo
+      prepend_if_ee 'Foo'
     end
     SOURCE
   end
@@ -110,7 +110,7 @@ describe RuboCop::Cop::InjectEnterpriseEditionModule do
   it 'does not flag including of regular modules' do
     expect_no_offenses(<<~SOURCE)
     class Foo
-      include Foo
+      include_if_ee 'Foo'
     end
     SOURCE
   end
@@ -118,51 +118,111 @@ describe RuboCop::Cop::InjectEnterpriseEditionModule do
   it 'does not flag extending using regular modules' do
     expect_no_offenses(<<~SOURCE)
     class Foo
-      extend Foo
+      extend_if_ee 'Foo'
     end
     SOURCE
   end
 
-  it 'does not flag the use of `prepend EE` on the last line' do
+  it 'does not flag the use of `prepend_if_ee EE` on the last line' do
     expect_no_offenses(<<~SOURCE)
     class Foo
     end
 
-    Foo.prepend(EE::Foo)
+    Foo.prepend_if_ee('EE::Foo')
     SOURCE
   end
 
-  it 'does not flag the use of `include EE` on the last line' do
+  it 'does not flag the use of `include_if_ee EE` on the last line' do
     expect_no_offenses(<<~SOURCE)
     class Foo
     end
 
-    Foo.include(EE::Foo)
+    Foo.include_if_ee('EE::Foo')
     SOURCE
   end
 
-  it 'does not flag the use of `extend EE` on the last line' do
+  it 'does not flag the use of `extend_if_ee EE` on the last line' do
     expect_no_offenses(<<~SOURCE)
     class Foo
     end
 
-    Foo.extend(EE::Foo)
+    Foo.extend_if_ee('EE::Foo')
     SOURCE
   end
 
   it 'autocorrects offenses by just disabling the Cop' do
     source = <<~SOURCE
     class Foo
-      prepend EE::Foo
-      include Bar
+      prepend_if_ee 'EE::Foo'
+      include_if_ee 'Bar'
     end
     SOURCE
 
     expect(autocorrect_source(source)).to eq(<<~SOURCE)
     class Foo
-      prepend EE::Foo # rubocop: disable Cop/InjectEnterpriseEditionModule
-      include Bar
+      prepend_if_ee 'EE::Foo' # rubocop: disable Cop/InjectEnterpriseEditionModule
+      include_if_ee 'Bar'
+    end
+    SOURCE
+  end
+
+  it 'disallows the use of prepend to inject an EE module' do
+    expect_offense(<<~SOURCE)
+    class Foo
+    end
+
+    Foo.prepend(EE::Foo)
+    ^^^^^^^^^^^^^^^^^^^^ EE modules must be injected using `include_if_ee`, `extend_if_ee`, or `prepend_if_ee`
+    SOURCE
+  end
+
+  it 'disallows the use of extend to inject an EE module' do
+    expect_offense(<<~SOURCE)
+    class Foo
     end
+
+    Foo.extend(EE::Foo)
+    ^^^^^^^^^^^^^^^^^^^ EE modules must be injected using `include_if_ee`, `extend_if_ee`, or `prepend_if_ee`
+    SOURCE
+  end
+
+  it 'disallows the use of include to inject an EE module' do
+    expect_offense(<<~SOURCE)
+    class Foo
+    end
+
+    Foo.include(EE::Foo)
+    ^^^^^^^^^^^^^^^^^^^^ EE modules must be injected using `include_if_ee`, `extend_if_ee`, or `prepend_if_ee`
+    SOURCE
+  end
+
+  it 'disallows the use of prepend_if_ee without a String' do
+    expect_offense(<<~SOURCE)
+    class Foo
+    end
+
+    Foo.prepend_if_ee(EE::Foo)
+                      ^^^^^^^ EE modules to inject must be specified as a String
+    SOURCE
+  end
+
+  it 'disallows the use of include_if_ee without a String' do
+    expect_offense(<<~SOURCE)
+    class Foo
+    end
+
+    Foo.include_if_ee(EE::Foo)
+                      ^^^^^^^ EE modules to inject must be specified as a String
+    SOURCE
+  end
+
+  it 'disallows the use of extend_if_ee without a String' do
+    expect_offense(<<~SOURCE)
+    class Foo
+    end
+
+    Foo.extend_if_ee(EE::Foo)
+                     ^^^^^^^ EE modules to inject must be specified as a String
     SOURCE
   end
 end