1 files changed, 15 insertions, 11 deletions
diff --git a/lib/tasks/gitlab/dev.rake b/lib/tasks/gitlab/dev.rake
index 5ee99dfc810..3117075b08b 100644
--- a/lib/tasks/gitlab/dev.rake
+++ b/lib/tasks/gitlab/dev.rake
@@ -1,18 +1,22 @@
 namespace :gitlab do
   namespace :dev do
     desc 'Checks if the branch would apply cleanly to EE'
-    task ee_compat_check: :environment do
-      return if defined?(Gitlab::License)
-      return unless ENV['CI']
+    task :ee_compat_check, [:branch] => :environment do |_, args|
+      opts =
+        if ENV['CI']
+          {
+            branch: ENV['CI_BUILD_REF_NAME'],
+            ce_repo: ENV['CI_BUILD_REPO']
+          }
+        else
+          unless args[:branch]
+            puts "Must specify a branch as an argument".color(:red)
+            exit 1
+          end
+          args
+        end
 
-      success =
-        Gitlab::EeCompatCheck.new(
-          branch: ENV['CI_BUILD_REF_NAME'],
-          check_dir: File.expand_path('ee-compat-check', __dir__),
-          ce_repo: ENV['CI_BUILD_REPO']
-        ).check
-
-      if success
+      if Gitlab::EeCompatCheck.new(opts || {}).check
         exit 0
       else
         exit 1