1
    2
    3
    4
    5
    6
    7
    8
    9
   10
   11
   12
   13
   14
   15
   16
   17
   18
   19
   20
   21
   22
   23
   24
   25
   26
   27
   28
   29
   30
   31
   32
   33
   34
   35
   36
   37
   38
   39
   40
   41
   42
   43
   44
   45
   46
   47
   48
   49
   50
   51
   52
   53
   54
   55
   56
   57
   58
   59
   60
   61
   62
   63
   64
   65
   66
   67
   68
   69
   70

content / test / gpu / bad_machine_finder / test_specs.py [blame]

# Copyright 2024 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Code for interacting with the test spec files in //testing/buildbot."""

import ast
import collections
import copy
import os
from typing import Any, Dict, Generator, List, Tuple

import gpu_path_util


class DimensionSet:
  """Represents a set of Swarming dimensions pulled from the test spec files.

  Handles the Swarming OR operator (|) and guarantees a consistent order when
  iterating over dimension key/value pairs.
  """

  def __init__(self, dimensions: Dict[str, str]):
    """
    Args:
      dimensions: A dict mapping dimension keys to values.
    """
    self._dimensions = collections.OrderedDict()
    for key in sorted(dimensions.keys()):
      value = dimensions[key]
      self._dimensions[key] = value.split('|')

  def Pairs(self) -> Generator[Tuple[str, List[str]], None, None]:
    """Iterates over the key/value pairs stored internally.

    Yields:
      (dimension_name, valid_values). |dimension_name| is a string containing
      the name of the dimension. |valid_values| is a list of one or more strings
      where each element is a valid value for |dimension_name|.
    """
    for dimension_name, valid_values in self._dimensions.items():
      yield dimension_name, valid_values

  def AsDict(self) -> Dict[str, List[str]]:
    """Returns a copy of the stored information as a dict."""
    return copy.deepcopy(self._dimensions)


def _LoadPylFile(filepath: str) -> Any:
  with open(filepath, encoding='utf-8') as infile:
    return ast.literal_eval(infile.read())


def GetMixinDimensions(mixin: str) -> 'DimensionSet':
  """Gets the dimensions provided by |mixin|.

  Args:
    mixin: The name of the mixin to look up.

  Returns:
    A dict mapping dimension names to values that |mixin| specifies.
  """
  mixin_content = _LoadPylFile(
      os.path.join(gpu_path_util.CHROMIUM_SRC_DIR, 'testing', 'buildbot',
                   'mixins.pyl'))
  dimensions = mixin_content.get(mixin, {}).get('swarming',
                                                {}).get('dimensions')
  if not dimensions:
    raise RuntimeError(
        f'Specified mixin {mixin} does not contain Swarming dimensions')
  return DimensionSet(dimensions)