External build system generator release 33

Update tools: yexport

1 files changed, 24 insertions, 6 deletions

diff --git a/build/scripts/split_unittest.py b/build/scripts/split_unittest.py
index 613bf507ad..8874b8b915 100644
--- a/build/scripts/split_unittest.py
+++ b/build/scripts/split_unittest.py
@@ -8,10 +8,28 @@ def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--split-factor", type=int, default=0)
     parser.add_argument("--shard", type=int, default=0)
+    parser.add_argument("--fork-mode", type=str, default="SEQUENTIAL")
     parser.add_argument("command", nargs=argparse.REMAINDER)
     return parser.parse_args()
 
 
+def get_sequential_chunk(tests, modulo, modulo_index):
+    chunk_size = len(tests) // modulo
+    not_used = len(tests) % modulo
+    shift = chunk_size + (modulo_index < not_used)
+    start = chunk_size * modulo_index + min(modulo_index, not_used)
+    end = start + shift
+    return [] if end > len(tests) else tests[start:end]
+
+
+def get_shuffled_chunk(tests, modulo, modulo_index):
+    result_tests = []
+    for i, test in enumerate(tests):
+        if i % modulo == modulo_index:
+            result_tests.append(test)
+    return result_tests
+
+
 def list_tests(binary):
     with tempfile.NamedTemporaryFile() as tmpfile:
         cmd = [binary, "--list-verbose", "--list-path", tmpfile.name]
@@ -27,12 +45,12 @@ def get_shard_tests(args):
     test_names = list_tests(args.command[0])
     test_names = sorted(test_names)
 
-    chunk_size = len(test_names) // args.split_factor
-    not_used = len(test_names) % args.split_factor
-    shift = chunk_size + (args.shard < not_used)
-    start = chunk_size * args.shard + min(args.shard, not_used)
-    end = start + shift
-    return [] if end > len(test_names) else test_names[start:end]
+    if args.fork_mode == "MODULO":
+        return get_shuffled_chunk(test_names, args.split_factor, args.shard)
+    elif args.fork_mode == "SEQUENTIAL":
+        return get_sequential_chunk(test_names, args.split_factor, args.shard)
+    else:
+        raise ValueError("detected unknown partition mode: {}".format(args.fork_mode))
 
 
 def get_shard_cmd_args(args):