Feat: --cutoff and "keep_workspace_files" options (#261)

Co-authored-by: merwanehamadi <merwanehamadi@gmail.com>

agbenchmark/agent_interface.py

@@ -38,6 +38,8 @@ def run_agent(
     timeout = cutoff
     if "--nc" in sys.argv:
         timeout = 100000
+    if "--cutoff" in sys.argv:
+        timeout = int(sys.argv[sys.argv.index("--cutoff") + 1])
 
     print(f"Running '{entry_path}' with timeout {timeout}")
 

agbenchmark/conftest.py

@@ -74,22 +74,24 @@ def workspace(config: Dict[str, Any]) -> Generator[str, None, None]:
 
     yield config["workspace"]
     # teardown after test function completes
-
-    for filename in os.listdir(output_path):
-        file_path = os.path.join(output_path, filename)
-        try:
-            if os.path.isfile(file_path) or os.path.islink(file_path):
-                os.unlink(file_path)
-            elif os.path.isdir(file_path):
-                shutil.rmtree(file_path)
-        except Exception as e:
-            print(f"Failed to delete {file_path}. Reason: {e}")
+    if not config.get("keep_workspace_files", False):
+        print("Emptying workspace")
+        for filename in os.listdir(output_path):
+            file_path = os.path.join(output_path, filename)
+            try:
+                if os.path.isfile(file_path) or os.path.islink(file_path):
+                    os.unlink(file_path)
+                elif os.path.isdir(file_path):
+                    shutil.rmtree(file_path)
+            except Exception as e:
+                print(f"Failed to delete {file_path}. Reason: {e}")
 
 
 def pytest_addoption(parser: Any) -> None:
     parser.addoption("--mock", action="store_true", default=False)
     parser.addoption("--category", action="store_true", default=False)
     parser.addoption("--nc", action="store_true", default=False)
+    parser.addoption("--cutoff", action="store_true", default=False)
     parser.addoption("--improve", action="store_true", default=False)
     parser.addoption("--maintain", action="store_true", default=False)
     parser.addoption("--test", action="store_true", default=None)

agbenchmark/start_benchmark.py

@@ -3,7 +3,7 @@ import os
 import sys
 from datetime import datetime
 from pathlib import Path
-from typing import Any
+from typing import Any, Optional
 
 import click
 import pytest
@@ -58,6 +58,7 @@ def cli() -> None:
     help="Run without dependencies (can be useful for a suite run)",
 )
 @click.option("--nc", is_flag=True, help="Run without cutoff")
+@click.option("--cutoff", default=None, help="Set or override tests cutoff (seconds)")
 def start(
     category: str,
     test: str,
@@ -67,6 +68,7 @@ def start(
     suite: str,
     no_dep: bool,
     nc: bool,
+    cutoff: Optional[int] = None,
 ) -> int:
     """Start the benchmark tests. If a category flag is provided, run the categories with that mark."""
     # Check if configuration file exists and is not empty
@@ -145,8 +147,18 @@ def start(
 
     if no_dep:
         pytest_args.append("--no_dep")
+
+    if nc and cutoff:
+        print(
+            "Error: You can't use both --nc and --cutoff at the same time. Please choose one."
+        )
+        return 1
+
     if nc:
         pytest_args.append("--nc")
+    if cutoff:
+        pytest_args.extend(["--cutoff", str(cutoff)])
+        print(f"Setting cuttoff override to {cutoff} seconds.")
 
     # when used as a library, the pytest directory to execute is in the CURRENT_DIRECTORY
     pytest_args.append(str(CURRENT_DIRECTORY))