From 82ce6f1e3ecb0bca1960d39870f7b8201dadc5f8 Mon Sep 17 00:00:00 2001
From: Rail Aliiev <rail@mozilla.com>
Date: Thu, 24 Sep 2015 17:10:11 -0400
Subject: [PATCH] Bug 1163998 - Balrog rule locking doesn't work in merge
 scripts. r=jlund DONTBUILD

---
 .../mozharness/mozilla/updates/balrog.py      | 31 +++++--------------
 .../scripts/merge_day/gecko_migration.py      |  7 ++++-
 2 files changed, 14 insertions(+), 24 deletions(-)

diff --git a/testing/mozharness/mozharness/mozilla/updates/balrog.py b/testing/mozharness/mozharness/mozilla/updates/balrog.py
index ab4085b80db9..4acbfa2db494 100644
--- a/testing/mozharness/mozharness/mozilla/updates/balrog.py
+++ b/testing/mozharness/mozharness/mozilla/updates/balrog.py
@@ -116,7 +116,8 @@ class BalrogMixin(object):
         c = self.config
         dirs = self.query_abs_dirs()
         submitter_script = os.path.join(
-            dirs["abs_tools_dir"], "scripts", "updates", "balrog-nightly-locker.py"
+            dirs["abs_tools_dir"], "scripts", "updates",
+            "balrog-nightly-locker.py"
         )
         credentials_file = os.path.join(
             dirs["base_work_dir"], c["balrog_credentials_file"]
@@ -126,6 +127,8 @@ class BalrogMixin(object):
             self.query_exe("python"),
             submitter_script,
             "--credentials-file", credentials_file,
+            "--api-root", c["balrog_api_root"],
+            "--username", c["balrog_username"],
         ]
         for r in rule_ids:
             cmd.extend(["-r", str(r)])
@@ -133,26 +136,8 @@ class BalrogMixin(object):
         if self._log_level_at_least(INFO):
             cmd.append("--verbose")
 
-        return_codes = []
-        for server in self.config["balrog_servers"]:
+        cmd.append("lock")
 
-            server_args = [
-                "--api-root", server["balrog_api_root"],
-                "--username", self._query_balrog_username(server)
-            ]
-
-            cmd.append("lock")
-
-            self.info("Calling Balrog rule locking script.")
-            return_code = self.retry(
-                self.run_command, attempts=5,
-                args=(cmd + server_args + ['lock'],),
-            )
-            if server["ignore_failures"]:
-                self.info("Ignoring result, ignore_failures set to True")
-            else:
-                return_codes.append(return_code)
-
-        # use the worst (max) code
-        if max(return_codes) != 0:
-            self.return_code = 1
+        self.info("Calling Balrog rule locking script.")
+        self.retry(self.run_command, attempts=5, args=cmd,
+                   kwargs={"halt_on_failure": True})
diff --git a/testing/mozharness/scripts/merge_day/gecko_migration.py b/testing/mozharness/scripts/merge_day/gecko_migration.py
index f02123548a8d..aa884ef289a1 100755
--- a/testing/mozharness/scripts/merge_day/gecko_migration.py
+++ b/testing/mozharness/scripts/merge_day/gecko_migration.py
@@ -47,11 +47,16 @@ class GeckoMigration(MercurialScript, BalrogMixin, VirtualenvMixin, SelfServeMix
             "default": "ffxbld <release@mozilla.com>",
             "help": "Specify what user to use to commit to hg.",
         }],
+        [['--balrog-api-root', ], {
+            "action": "store",
+            "dest": "balrog_api_root",
+            "type": "string",
+            "help": "Specify Balrog API root URL.",
+        }],
         [['--balrog-username', ], {
             "action": "store",
             "dest": "balrog_username",
             "type": "string",
-            "default": "ffxbld",
             "help": "Specify what user to connect to Balrog with.",
         }],
         [['--balrog-credentials-file', ], {