From 9b6cbfe7aec695d7a0c788e853c69f1cdff2c7eb Mon Sep 17 00:00:00 2001
From: Yingchi Long <longyingchi24s@ict.ac.cn>
Date: Tue, 22 Apr 2025 22:52:37 +0800
Subject: [PATCH] speccpu: add `SPEC` class

---
 src/speccpu/__init__.py | 33 ++++++++++++++++++++++++++++
 test/test_spec.py       | 48 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 80 insertions(+), 1 deletion(-)

diff --git a/src/speccpu/__init__.py b/src/speccpu/__init__.py
index 9fbdb2c..0f6afd1 100644
--- a/src/speccpu/__init__.py
+++ b/src/speccpu/__init__.py
@@ -84,3 +84,36 @@ def find_build(build_dir: Path) -> Path:
         if entry.startswith("build"):
             return build_dir / entry
     raise RuntimeError(f"SPEC build directory not found at {build_dir}")
+
+
+class SPEC:
+    dir: Path
+
+    def __init__(self, dir: Path):
+        self.dir = dir
+
+    @property
+    def config_dir(self) -> Path:
+        return self.dir / "config"
+
+    def env(self):
+        return create_spec_env(self.dir)
+
+    def mkcmd_runcpu(
+        self,
+        config: str,
+        benchmarks: list[str],
+        setprocgroup: bool = True,
+        workload: str = "ref",
+        output_root: Path | None = None,
+    ):
+        return [
+            "runcpu",
+            *(["--setprocgroup"] if setprocgroup else []),
+            "-i",
+            workload,
+            "-c",
+            config,
+            *(["--outputdir", str(output_root)] if output_root else []),
+            *benchmarks,
+        ]
diff --git a/test/test_spec.py b/test/test_spec.py
index 919fc36..871c518 100644
--- a/test/test_spec.py
+++ b/test/test_spec.py
@@ -1,8 +1,9 @@
 from pathlib import Path
+from unittest import mock
 
 import pytest
 
-from speccpu import find_build
+from speccpu import SPEC, find_build
 
 
 class TestFindBuild:
@@ -33,3 +34,48 @@ class TestFindBuild:
         build_parent = Path(tmpdir)
         with pytest.raises(RuntimeError):
             find_build(build_parent)
+
+
+class TestSPEC:
+    def test_init(self):
+        spec_dir = Path("/path/to/spec")
+        spec = SPEC(spec_dir)
+        assert spec.dir == spec_dir
+
+    def test_config_dir(self):
+        spec_dir = Path("/path/to/spec")
+        spec = SPEC(spec_dir)
+        assert spec.config_dir == spec_dir / "config"
+
+    def test_env(self):
+        spec_dir = Path("/path/to/spec")
+        spec = SPEC(spec_dir)
+        with mock.patch("speccpu.create_spec_env") as mock_create_env:
+            spec.env()
+            mock_create_env.assert_called_once_with(spec_dir)
+
+    def test_mkcmd_runcpu_default(self):
+        spec = SPEC(Path("/path/to/spec"))
+        cmd = spec.mkcmd_runcpu("myconfig", ["benchmark1", "benchmark2"])
+        assert cmd == [
+            "runcpu",
+            "--setprocgroup",
+            "-i",
+            "ref",
+            "-c",
+            "myconfig",
+            "benchmark1",
+            "benchmark2",
+        ]
+
+    def test_mkcmd_runcpu_no_procgroup(self):
+        spec = SPEC(Path("/path/to/spec"))
+        cmd = spec.mkcmd_runcpu("myconfig", ["benchmark1"], setprocgroup=False)
+        assert "--setprocgroup" not in cmd
+
+    def test_mkcmd_runcpu_with_output_dir(self):
+        spec = SPEC(Path("/path/to/spec"))
+        output_dir = Path("/output/dir")
+        cmd = spec.mkcmd_runcpu("myconfig", ["benchmark1"], output_root=output_dir)
+        assert "--outputdir" in cmd
+        assert str(output_dir) in cmd