1 files changed, 29 insertions, 6 deletions
diff --git a/config/initializers/0_inject_enterprise_edition_module.rb b/config/initializers/0_inject_enterprise_edition_module.rb
index b3ebb44ef25..7478727f869 100644
--- a/config/initializers/0_inject_enterprise_edition_module.rb
+++ b/config/initializers/0_inject_enterprise_edition_module.rb
@@ -6,12 +6,7 @@ module InjectEnterpriseEditionModule
   def prepend_if_ee(constant, with_descendants: false)
     return unless Gitlab.ee?
 
-    ee_module = constant.constantize
-    prepend(ee_module)
-
-    if with_descendants
-      descendants.each { |descendant| descendant.prepend(ee_module) }
-    end
+    prepend_module(constant.constantize, with_descendants)
   end
 
   def extend_if_ee(constant)
@@ -21,6 +16,34 @@ module InjectEnterpriseEditionModule
   def include_if_ee(constant)
     include(constant.constantize) if Gitlab.ee?
   end
+
+  def prepend_ee_mod(with_descendants: false)
+    return unless Gitlab.ee?
+
+    prepend_module(ee_module, with_descendants)
+  end
+
+  def extend_ee_mod
+    extend(ee_module) if Gitlab.ee?
+  end
+
+  def include_ee_mod
+    include(ee_module) if Gitlab.ee?
+  end
+
+  private
+
+  def prepend_module(mod, with_descendants)
+    prepend(mod)
+
+    if with_descendants
+      descendants.each { |descendant| descendant.prepend(mod) }
+    end
+  end
+
+  def ee_module
+    ::EE.const_get(name, false)
+  end
 end
 
 Module.prepend(InjectEnterpriseEditionModule)