amend! [unittest] Add option to allow disabling sharding in unittest

Use for f-strings for error messages. Use common env key/values.

Author: zeroomega

llvm/utils/lit/lit/formats/googletest.py

@@ -147,6 +147,10 @@ def execute(self, test, litConfig):
         from lit.cl_arguments import TestOrder
 
         use_shuffle = TestOrder(litConfig.order) == TestOrder.RANDOM
+        shard_env = {
+            "GTEST_OUTPUT": "json:" + test.gtest_json_file,
+            "GTEST_SHUFFLE": "1" if use_shuffle else "0",
+        }
         if not litConfig.disableGTestSharding:
             testPath, testName = os.path.split(test.getSourcePath())
             while not os.path.exists(testPath):
@@ -157,19 +161,14 @@ def execute(self, test, litConfig):
 
             testName, total_shards = os.path.split(testName)
             testName, shard_idx = os.path.split(testName)
-            shard_env = {
-                "GTEST_OUTPUT": "json:" + test.gtest_json_file,
-                "GTEST_SHUFFLE": "1" if use_shuffle else "0",
-                "GTEST_TOTAL_SHARDS": os.environ.get(
-                    "GTEST_TOTAL_SHARDS", total_shards
-                ),
-                "GTEST_SHARD_INDEX": os.environ.get("GTEST_SHARD_INDEX", shard_idx),
-            }
-        else:
-            shard_env = {
-                "GTEST_OUTPUT": "json:" + test.gtest_json_file,
-                "GTEST_SHUFFLE": "1" if use_shuffle else "0",
-            }
+            shard_env.update(
+                {
+                    "GTEST_TOTAL_SHARDS": os.environ.get(
+                        "GTEST_TOTAL_SHARDS", total_shards
+                    ),
+                    "GTEST_SHARD_INDEX": os.environ.get("GTEST_SHARD_INDEX", shard_idx),
+                }
+            )
         test.config.environment.update(shard_env)
 
         cmd = [testPath]

llvm/utils/lit/tests/Inputs/googletest-no-sharding/DummySubDir/OneTest.py

@@ -5,7 +5,7 @@
 
 if len(sys.argv) == 3 and sys.argv[1] == "--gtest_list_tests":
     if sys.argv[2] != "--gtest_filter=-*DISABLED_*":
-        raise ValueError("unexpected argument: %s" % (sys.argv[2]))
+        raise ValueError(f"unexpected argument: {sys.argv[2]}")
     print(
         """\
 FirstTest.
@@ -21,14 +21,14 @@
     sys.exit(0)
 elif len(sys.argv) != 1:
     # sharding and json output are specified using environment variables
-    raise ValueError("unexpected argument: %r" % (" ".join(sys.argv[1:])))
+    raise ValueError(f"unexpected argument: {' '.join(sys.argv[1:])!r}")
 
 for e in ["GTEST_OUTPUT"]:
     if e not in os.environ:
-        raise ValueError("missing environment variables: " + e)
+        raise ValueError(f"missing environment variables: {e}")
 
 if not os.environ["GTEST_OUTPUT"].startswith("json:"):
-    raise ValueError("must emit json output: " + os.environ["GTEST_OUTPUT"])
+    raise ValueError(f"must emit json output: {os.environ['GTEST_OUTPUT']}")
 
 output = """\
 {