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.

1 files changed, 41 insertions, 7 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