# Copyright 2020 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
import collections
import logging
import os
from typing import Tuple, List
import test_runner_errors
LOGGER = logging.getLogger(__name__)
class ShardingError(test_runner_errors.Error):
"""Error related with sharding logic."""
pass
class ExcessShardsError(ShardingError):
"""The test module is misconfigured to have more shards than test cases"""
def __init__(self, num_shards, num_test_cases):
super(ExcessShardsError, self).__init__(
f'The test module is misconfigured to have more shards than test cases.'
f' Shards: {num_shards} Test Cases: {num_test_cases}')
def gtest_shard_index():
"""Returns shard index in environment, or 0 if not in sharding environment."""
return int(os.getenv('GTEST_SHARD_INDEX', 0))
def gtest_total_shards():
"""Returns total shard count in environment, or 1 if not in environment."""
return int(os.getenv('GTEST_TOTAL_SHARDS', 1))
def balance_into_sublists(test_counts: collections.Counter,
total_shards: int) -> List[List[str]]:
"""Augment the result of otool into balanced sublists
Args:
test_counts: (collections.Counter) dict of test_case to test case numbers
total_shards: (int) total number of shards this was divided into
Returns:
list of list of test classes
"""
class Shard(object):
"""Stores list of test classes and number of all tests"""
def __init__(self):
self.test_classes = []
self.size = 0
shards = [Shard() for i in range(total_shards)]
# Balances test classes between shards to have
# approximately equal number of tests per shard.
for test_class, number_of_test_methods in test_counts.most_common():
min_shard = min(shards, key=lambda shard: shard.size)
min_shard.test_classes.append(test_class)
min_shard.size += number_of_test_methods
LOGGER.debug('%s test case is allocated to shard %s with %s test methods' %
(test_class, shards.index(min_shard), number_of_test_methods))
sublists = [shard.test_classes for shard in shards]
return sublists
def shard_eg_test_cases(all_eg_test_names: List[Tuple[str, str]]) -> List[str]:
"""Shard test cases into total_shards, and determine which test cases to
run for this shard.
Raises:
ExcessShardsError: If there exist more shards than test_cases
Args:
all_eg_test_names: A list of all EG test methods present in the
-Runner.app binary. Each list element is a tuple in the form
(test_case, test_method)
Returns: a list of test cases to execute on this shard
"""
shard_index = gtest_shard_index()
total_shards = gtest_total_shards()
test_counts = collections.Counter(
test_class for test_class, _ in all_eg_test_names)
# Ensure shard and total shard is int
shard_index = int(shard_index)
total_shards = int(total_shards)
total_test_cases = len(test_counts)
if total_shards > total_test_cases:
raise ExcessShardsError(total_shards, total_test_cases)
sublists = balance_into_sublists(test_counts, total_shards)
tests = sublists[shard_index]
LOGGER.info("Tests to be executed this round: {}".format(tests))
return tests
Codebase Browser
chromium