refactor: with sonnet

Author: paulz

src/cat_ai/runner.py

@@ -1,25 +1,59 @@
 import os
+from typing import Callable, List, Optional
 
 from .reporter import Reporter
-from typing import Callable, Any
 
 
 class Runner:
-    def __init__(self, test_function: Callable[..., Any], reporter: Reporter) -> None:
+    """Executes test functions and collects results using a reporter."""
+
+    def __init__(self, test_function: Callable[..., bool], reporter: Reporter) -> None:
+        """
+        Initialize the Runner with a test function and reporter.
+
+        Args:
+            test_function: Function to execute during test runs
+            reporter: Reporter instance to track and report test results
+        """
         self.reporter = reporter
         self.test_function = test_function
 
     @staticmethod
-    def sample_size(default_size: int = 1) -> int:
+    def get_sample_size(default_size: int = 1) -> int:
+        """
+        Get sample size from environment variable or use default.
+
+        Args:
+            default_size: Default sample size if not specified in environment
+
+        Returns:
+            Number of test runs to perform
+        """
         return int(os.getenv("CAT_AI_SAMPLE_SIZE", str(default_size)))
 
-    def run_once(self, run_number: int = 0) -> Any:
+    def run_once(self, run_number: int = 0) -> bool:
+        """
+        Execute the test function once.
+
+        Args:
+            run_number: Current run index for reporting
+
+        Returns:
+            Result from the test function
+        """
         self.reporter.run_number = run_number
-        result = self.test_function(reporter=self.reporter)
-        return result
-
-    def run_loop(self, tries: int = sample_size()) -> list[Any]:
-        results = []
-        for x in range(0, tries):
-            results.append(self.run_once(x))
-        return results
+        return self.test_function(reporter=self.reporter)
+
+    def run(self, sample_size: Optional[int] = None) -> List[bool]:
+        """
+        Execute the test function multiple times based on sample size.
+
+        Args:
+            sample_size: Number of times to run the test, defaults to
+                         value from get_sample_size() if None
+
+        Returns:
+            List of results from all test runs
+        """
+        runs = sample_size or self.get_sample_size()
+        return [self.run_once(i) for i in range(runs)]

tests/cat_ai/test_runner.py

@@ -1,6 +1,5 @@
-import pytest
-from src.cat_ai.runner import Runner
 from src.cat_ai.reporter import Reporter
+from src.cat_ai.runner import Runner
 
 
 # Dummy test function that will be passed to Runner
@@ -12,11 +11,11 @@ def dummy_test_function(reporter: Reporter):
 def test_runner_sample_size(monkeypatch):
     # Set an environment variable to test
     monkeypatch.setenv("CAT_AI_SAMPLE_SIZE", "5")
-    assert Runner.sample_size() == 5
+    assert Runner.get_sample_size() == 5
 
     # Test default size
     monkeypatch.delenv("CAT_AI_SAMPLE_SIZE", raising=False)
-    assert Runner.sample_size(default_size=3) == 3
+    assert Runner.get_sample_size(default_size=3) == 3
 
 
 def test_run_once():
@@ -32,22 +31,38 @@ def test_run_once():
     assert reporter.run_number == 1
 
 
-def test_run_loop(monkeypatch):
-    # Set the environment variable for a controlled loop
+def test_run():
+    # Create a Reporter with necessary arguments
+    reporter = Reporter(test_name="test_run", output_dir="/tmp")
+
+    # Initialize Runner with dummy test function and Reporter
+    runner = Runner(test_function=dummy_test_function, reporter=reporter)
+
+    # Test with explicit sample size parameter
+    results = runner.run(sample_size=3)
+    assert len(results) == 3
+    assert all(res.startswith("Running test with run number ") for res in results)
+    expected_results = [
+        "Running test with run number 0",
+        "Running test with run number 1",
+        "Running test with run number 2",
+    ]
+    assert results == expected_results
+
+
+def test_run_with_env_variable(monkeypatch):
+    # Set the environment variable for a controlled test
     monkeypatch.setenv("CAT_AI_SAMPLE_SIZE", "3")
 
     # Create a Reporter with necessary arguments
-    reporter = Reporter(test_name="test_run_loop", output_dir="/tmp")
+    reporter = Reporter(test_name="test_run_with_env", output_dir="/tmp")
 
     # Initialize Runner with dummy test function and Reporter
     runner = Runner(test_function=dummy_test_function, reporter=reporter)
 
-    # Test run_loop
-    tries = Runner.sample_size()
-
-    results = runner.run_loop(tries=tries)
+    # Test without explicit sample size (should use environment variable)
+    results = runner.run()
     assert len(results) == 3
-    assert all(res.startswith("Running test with run number ") for res in results)
     expected_results = [
         "Running test with run number 0",
         "Running test with run number 1",