cbuildbot: add HWTest DUT dims override flag

BUG=b:188802911
TEST=unit tests

Change-Id: Ia1ce51c4c6ebd1c958dee5973c331d1344c0a4c2
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/2992014
Reviewed-by: LaMont Jones <lamontjones@chromium.org>
Reviewed-by: Dhanya Ganesh <dhanyaganesh@chromium.org>
Commit-Queue: Jared Loucks <jaredloucks@google.com>
Tested-by: Jared Loucks <jaredloucks@google.com>
diff --git a/scripts/cbuildbot.py b/scripts/cbuildbot.py
index a13ba0c..ebed54d 100644
--- a/scripts/cbuildbot.py
+++ b/scripts/cbuildbot.py
@@ -21,6 +21,7 @@
 from chromite.cbuildbot import repository
 from chromite.cbuildbot import topology
 from chromite.cbuildbot.stages import completion_stages
+from chromite.cbuildbot.stages import test_stages
 from chromite.lib import builder_status_lib
 from chromite.lib import cidb
 from chromite.lib import cgroups
@@ -49,6 +50,9 @@
 _DEFAULT_INT_BUILDROOT = 'trybot-internal'
 _BUILDBOT_REQUIRED_BINARIES = ('pbzip2',)
 _API_VERSION_ATTR = 'api_version'
+BOARD_DIM_LABEL = 'label-board'
+MODEL_DIM_LABEL = 'label-model'
+POOL_DIM_LABEL = 'label-pool'
 
 
 def _BackupPreviousLog(log_file, backup_limit=25):
@@ -311,6 +315,13 @@
       'Options used to configure tryjob behavior.')
   group.add_remote_option('--hwtest', action='store_true', default=False,
                           help='Run the HWTest stage (tests on real hardware)')
+  group.add_option('--hwtest_dut_dimensions', type='string',
+                   action='split_extend', default=None,
+                   help='Space-separated list of key:val Swarming bot '
+                        'dimensions to run each builders SkylabHWTest '
+                        'stages against (this overrides the configured '
+                        'DUT dimensions for each test). Requires at least '
+                        '"label-board", "label-model", and "label-pool".')
   group.add_remote_option('--channel', action='split_extend', dest='channels',
                           default=[],
                           help='Specify a channel for a payloads trybot. Can '
@@ -623,6 +634,36 @@
   # We force --debug to be set for builds that are not 'official'.
   options.debug = options.debug or not options.buildbot
 
+  options.hwtest_dut_override = ParseHWTestDUTDims(
+    options.hwtest_dut_dimensions)
+
+def ParseHWTestDUTDims(dims):
+  """Parse HWTest DUT dimensions into a valid HWTestDUTOverride object.
+
+  Raises an error if any of board, model, or pool is missing.
+  """
+  if not dims:
+    return None
+  board = model = pool = None
+  extra_dims = []
+  for dim in dims:
+    if dim.startswith(BOARD_DIM_LABEL):
+      # Remove one extra character to account for the ":" or "=" symbol
+      # separating the label from the dimension itself.
+      board = dim[len(BOARD_DIM_LABEL)+1:]
+    elif dim.startswith(MODEL_DIM_LABEL):
+      model = dim[len(MODEL_DIM_LABEL)+1:]
+    elif dim.startswith(POOL_DIM_LABEL):
+      pool = dim[len(POOL_DIM_LABEL)+1:]
+    else:
+      extra_dims.append(dim)
+
+  if not (board and model and pool):
+    cros_build_lib.Die('HWTest DUT dimensions must include board, model, and '
+                       'pool (given %s).' % dims)
+
+  return test_stages.HWTestDUTOverride(board, model, pool, extra_dims)
+
 
 # pylint: disable=unused-argument
 def _PostParseCheck(parser, options, site_config):