Bug 1523444 - Don't hardcode branch names in partials.py helper, let it inherit from scriptworkers. r=sfraser

Differential Revision: https://phabricator.services.mozilla.com/D18233

--HG--
extra : moz-landing-system : lando

taskcluster/taskgraph/util/partials.py

@@ -8,6 +8,9 @@ import requests
 import redo
 
 import logging
+
+from taskgraph.util.scriptworker import BALROG_SCOPE_ALIAS_TO_PROJECT, BALROG_SERVER_SCOPES
+
 logger = logging.getLogger(__name__)
 
 PLATFORM_RENAMES = {
@@ -170,10 +173,20 @@ def get_release_builds(release, branch):
 
 
 def _get_balrog_api_root(branch):
-    if branch in ('mozilla-central', 'mozilla-beta', 'mozilla-release') or 'mozilla-esr' in branch:
-        return 'https://aus5.mozilla.org/api/v1'
+    # Query into the scopes scriptworker uses to make sure we check against the same balrog server
+    # That our jobs would use.
+    scope = None
+    for alias, projects in BALROG_SCOPE_ALIAS_TO_PROJECT:
+        if branch in projects and alias in BALROG_SERVER_SCOPES:
+            scope = BALROG_SERVER_SCOPES[alias]
+            break
     else:
+        scope = BALROG_SERVER_SCOPES['default']
+
+    if scope == u'balrog:server:dep':
         return 'https://aus5.stage.mozaws.net/api/v1'
+    else:
+        return 'https://aus5.mozilla.org/api/v1'
 
 
 def find_localtest(fileUrls):