From 93d18bc3bdd928d23a6786aff6f50e0627287240 Mon Sep 17 00:00:00 2001
From: papertager <2567587994@qq.com>
Date: Tue, 9 Jun 2026 22:09:39 +0800
Subject: [PATCH] Add mcoplib test shard planner

---
 tools/plan_test_shards.py          | 47 ++++++++++++++++++++++++++++++
 unit_test/test_plan_test_shards.py | 14 +++++++++
 2 files changed, 61 insertions(+)
 create mode 100644 tools/plan_test_shards.py
 create mode 100644 unit_test/test_plan_test_shards.py

diff --git a/tools/plan_test_shards.py b/tools/plan_test_shards.py
new file mode 100644
index 0000000..f51b231
--- /dev/null
+++ b/tools/plan_test_shards.py
@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+"""Create deterministic shards for mcoplib unit tests."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+
+def discover_tests(root: Path) -> list[str]:
+    return sorted(path.relative_to(root).as_posix() for path in (root / "unit_test").glob("test_*.py"))
+
+
+def plan_shards(tests: list[str], shard_count: int) -> list[list[str]]:
+    if shard_count <= 0:
+        raise ValueError("shard_count must be positive")
+    shards = [[] for _ in range(shard_count)]
+    for index, test in enumerate(tests):
+        shards[index % shard_count].append(test)
+    return shards
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--root", type=Path, default=Path.cwd(), help="repository root")
+    parser.add_argument("--shards", type=int, default=4, help="number of shards to create")
+    parser.add_argument("--output", type=Path, help="write JSON plan to this path")
+    args = parser.parse_args()
+
+    tests = discover_tests(args.root)
+    shards = plan_shards(tests, args.shards)
+    payload = {
+        "test_count": len(tests),
+        "shard_count": args.shards,
+        "shards": [{"index": index, "tests": shard} for index, shard in enumerate(shards)],
+    }
+    text = json.dumps(payload, indent=2, ensure_ascii=False)
+    if args.output:
+        args.output.write_text(text + "\n", encoding="utf-8")
+    else:
+        print(text)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/unit_test/test_plan_test_shards.py b/unit_test/test_plan_test_shards.py
new file mode 100644
index 0000000..fe81832
--- /dev/null
+++ b/unit_test/test_plan_test_shards.py
@@ -0,0 +1,14 @@
+import unittest
+
+from tools.plan_test_shards import plan_shards
+
+
+class PlanTestShardsTest(unittest.TestCase):
+    def test_distributes_tests_deterministically(self):
+        shards = plan_shards(["a.py", "b.py", "c.py", "d.py", "e.py"], 2)
+
+        self.assertEqual(shards, [["a.py", "c.py", "e.py"], ["b.py", "d.py"]])
+
+
+if __name__ == "__main__":
+    unittest.main()