Add infra recipe module, use for updating Go DEPS

BUG=skia:5879
GOLD_TRYBOT_URL= https://gold.skia.org/search?issue=2444883002

Review-Url: https://codereview.chromium.org/2444883002

6 files changed, 73 insertions, 22 deletions

diff --git a/infra/bots/recipe_modules/infra/__init__.py b/infra/bots/recipe_modules/infra/__init__.py
new file mode 100644
index 0000000000..ba20b52f95
--- /dev/null
+++ b/infra/bots/recipe_modules/infra/__init__.py
@@ -0,0 +1,9 @@
+# Copyright 2016 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+DEPS = [
+  'recipe_engine/step',
+  'run',
+  'vars',
+]
diff --git a/infra/bots/recipe_modules/infra/api.py b/infra/bots/recipe_modules/infra/api.py
new file mode 100644
index 0000000000..17d3729b87
--- /dev/null
+++ b/infra/bots/recipe_modules/infra/api.py
@@ -0,0 +1,32 @@
+# Copyright 2016 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+from recipe_engine import recipe_api
+
+
+INFRA_GO_PKG = 'go.skia.org/infra'
+UPDATE_GO_ATTEMPTS = 5
+
+
+class InfraApi(recipe_api.RecipeApi):
+  @property
+  def go_env(self):
+    return {'GOPATH': self.gopath}
+
+  @property
+  def gopath(self):
+    return self.m.vars.checkout_root.join('gopath')
+
+  def update_go_deps(self):
+    """Attempt to update go dependencies.
+
+    This fails flakily sometimes, so perform multiple attempts.
+    """
+    self.m.run.with_retry(
+        self.m.step,
+        'update go pkgs',
+        UPDATE_GO_ATTEMPTS,
+        cmd=['go', 'get', '-u', '%s/...' % INFRA_GO_PKG],
+        env=self.go_env)
diff --git a/infra/bots/recipe_modules/run/api.py b/infra/bots/recipe_modules/run/api.py
index ffcced1947..3ef69d806b 100644
--- a/infra/bots/recipe_modules/run/api.py
+++ b/infra/bots/recipe_modules/run/api.py
@@ -115,3 +115,15 @@ for pattern in build_products_whitelist:
 ''' % str(BUILD_PRODUCTS_ISOLATE_WHITELIST),
         args=[src, dst],
         infra_step=True)
+
+  def with_retry(self, steptype, name, attempts, *args, **kwargs):
+    for attempt in xrange(attempts):
+      step_name = name
+      if attempt > 0:
+        step_name += ' (attempt %d)' % (attempt + 1)
+      try:
+        steptype(step_name, *args, **kwargs)
+        return
+      except self.m.step.StepFailure:
+        if attempt == attempts - 1:
+          raise
diff --git a/infra/bots/recipes/swarm_RecreateSKPs.expected/Housekeeper-Weekly-RecreateSKPs.json b/infra/bots/recipes/swarm_RecreateSKPs.expected/Housekeeper-Weekly-RecreateSKPs.json
index 05e83a393c..15feac8777 100644
--- a/infra/bots/recipes/swarm_RecreateSKPs.expected/Housekeeper-Weekly-RecreateSKPs.json
+++ b/infra/bots/recipes/swarm_RecreateSKPs.expected/Housekeeper-Weekly-RecreateSKPs.json
@@ -245,6 +245,18 @@
   },
   {
     "cmd": [
+      "go",
+      "get",
+      "-u",
+      "go.skia.org/infra/..."
+    ],
+    "env": {
+      "GOPATH": "[CUSTOM_/_B_WORK]/gopath"
+    },
+    "name": "update go pkgs"
+  },
+  {
+    "cmd": [
       "python",
       "-u",
       "\nimport os\nimport urllib2\n\nTOKEN_FILE = '.depot_tools_oauth2_tokens'\nTOKEN_FILE_BACKUP = '.depot_tools_oauth2_tokens.old'\nTOKEN_URL = 'http://metadata/computeMetadata/v1/project/attributes/depot_tools_auth_update_skps'\n\nreq = urllib2.Request(TOKEN_URL, headers={'Metadata-Flavor': 'Google'})\ncontents = urllib2.urlopen(req).read()\n\nhome = os.path.expanduser('~')\ntoken_file = os.path.join(home, TOKEN_FILE)\nif os.path.isfile(token_file):\n  os.rename(token_file, os.path.join(home, TOKEN_FILE_BACKUP))\n\nwith open(token_file, 'w') as f:\n  f.write(contents)\n"
@@ -282,6 +294,7 @@
     "cwd": "[CUSTOM_/_B_WORK]/skia",
     "env": {
       "CHROME_HEADLESS": "1",
+      "GOPATH": "[CUSTOM_/_B_WORK]/gopath",
       "PATH": "[DEPOT_TOOLS]:%(PATH)s"
     },
     "name": "Upload SKPs"
diff --git a/infra/bots/recipes/swarm_RecreateSKPs.py b/infra/bots/recipes/swarm_RecreateSKPs.py
index eb92484068..d13e7507cc 100644
--- a/infra/bots/recipes/swarm_RecreateSKPs.py
+++ b/infra/bots/recipes/swarm_RecreateSKPs.py
@@ -10,6 +10,7 @@ DEPS = [
   'build/file',
   'core',
   'depot_tools/gclient',
+  'infra',
   'recipe_engine/path',
   'recipe_engine/properties',
   'recipe_engine/python',
@@ -136,9 +137,11 @@ def RunSteps(api):
 
   # Upload the SKPs.
   if 'Canary' not in api.properties['buildername']:
+    api.infra.update_go_deps()
     cmd = ['python',
            api.vars.skia_dir.join('infra', 'bots', 'upload_skps.py'),
            '--target_dir', output_dir]
+    env.update(api.infra.go_env)
     with depot_tools_auth(api, UPDATE_SKPS_KEY):
       api.step('Upload SKPs',
                cmd=cmd,
diff --git a/infra/bots/recipes/swarm_infra.py b/infra/bots/recipes/swarm_infra.py
index bfc004b138..a790fd5f59 100644
--- a/infra/bots/recipes/swarm_infra.py
+++ b/infra/bots/recipes/swarm_infra.py
@@ -8,37 +8,19 @@
 
 DEPS = [
   'core',
+  'infra',
   'recipe_engine/path',
   'recipe_engine/properties',
   'recipe_engine/step',
+  'run',
   'vars',
 ]
 
 
-UPDATE_GO_ATTEMPTS = 5
-
-
 def RunSteps(api):
   api.vars.setup()
   api.core.checkout_steps()
-
-  # Attempt to update go dependencies. This fails flakily sometimes, so perform
-  # multiple attempts.
-  gopath = api.vars.checkout_root.join('gopath')
-  env = {'GOPATH': gopath}
-  name = 'update go pkgs'
-  for attempt in xrange(UPDATE_GO_ATTEMPTS):
-    step_name = name
-    if attempt > 0:
-      step_name += ' (attempt %d)' % (attempt + 1)
-    try:
-      api.step(step_name,
-               cmd=['go', 'get', '-u', 'go.skia.org/infra/...'],
-               env=env)
-      break
-    except api.step.StepFailure:
-      if attempt == UPDATE_GO_ATTEMPTS - 1:
-        raise
+  api.infra.update_go_deps()
 
   # Run the infra tests.
   infra_tests = api.vars.skia_dir.join(
@@ -46,7 +28,7 @@ def RunSteps(api):
   api.step('infra_tests',
            cmd=['python', infra_tests],
            cwd=api.vars.skia_dir,
-           env=env)
+           env=api.infra.go_env)
 
 
 def GenTests(api):