upgraded testing

Author: cemde

maseval/core/registry.py

@@ -319,9 +319,14 @@ def collect_usage(self) -> Dict[str, Any]:
                         usage[category] = {}
                     usage[category][comp_name] = usage_dict
 
-                # Accumulate into persistent aggregates (thread-safe)
+                # Accumulate into persistent aggregates (thread-safe).
+                # _usage_total starts as Usage(cost=None); adding to it would
+                # poison the cost (None + X = None).  Assign directly on first use.
                 with self._usage_lock:
-                    self._usage_total = self._usage_total + component_usage
+                    if self._usage_total.cost is None and not self._usage_total.units:
+                        self._usage_total = component_usage
+                    else:
+                        self._usage_total = self._usage_total + component_usage
                     if key in self._usage_by_component:
                         self._usage_by_component[key] = self._usage_by_component[key] + component_usage
                     else:

tests/test_core/test_benchmark/test_usage_collection.py

@@ -0,0 +1,97 @@
+"""Test usage collection through the benchmark execution loop.
+
+These tests verify that benchmark.run() collects usage from registered
+model adapters and includes it in report dicts.
+"""
+
+import pytest
+from maseval import TaskQueue
+from maseval.core.usage import StaticPricingCalculator
+
+
+@pytest.mark.core
+class TestBenchmarkUsageCollection:
+    """Tests for usage collection during benchmark execution."""
+
+    def test_usage_in_report(self):
+        """Benchmark run includes a 'usage' key in each report."""
+        from conftest import DummyBenchmark
+
+        tasks = TaskQueue.from_list([{"query": "Test", "environment_data": {}}])
+        benchmark = DummyBenchmark()
+
+        reports = benchmark.run(tasks, agent_data={"model": "test"})
+
+        assert "usage" in reports[0]
+        usage = reports[0]["usage"]
+        assert "metadata" in usage
+        assert "models" in usage
+        assert "agents" in usage
+
+    def test_usage_has_correct_structure(self):
+        """Usage dict has the expected category keys and metadata."""
+        from conftest import DummyBenchmark
+
+        tasks = TaskQueue.from_list([{"query": "Test", "environment_data": {}}])
+        benchmark = DummyBenchmark()
+
+        reports = benchmark.run(tasks, agent_data={"model": "test"})
+
+        usage = reports[0]["usage"]
+        assert "metadata" in usage
+        assert "total_components" in usage["metadata"]
+        assert "timestamp" in usage["metadata"]
+
+    def test_model_with_usage_appears_in_report(self):
+        """A model adapter that reports usage has its tokens in the report."""
+        from conftest import DummyModelAdapter, DummyBenchmark
+
+        class UsageBenchmark(DummyBenchmark):
+            def get_model_adapter(self, model_id, **kwargs):
+                return DummyModelAdapter(
+                    model_id=model_id,
+                    usage={
+                        "input_tokens": 100,
+                        "output_tokens": 50,
+                        "total_tokens": 150,
+                    },
+                )
+
+        tasks = TaskQueue.from_list([{"query": "Test", "environment_data": {}}])
+        benchmark = UsageBenchmark()
+
+        reports = benchmark.run(tasks, agent_data={"model": "test"})
+
+        # The DummyBenchmark doesn't register a model via register(), so
+        # the model's usage won't appear unless the benchmark hooks it up.
+        # This test verifies the usage structure exists.
+        assert "usage" in reports[0]
+
+    def test_usage_persists_across_task_repetitions(self):
+        """Benchmark.usage accumulates across multiple tasks."""
+        from conftest import DummyBenchmark
+
+        tasks = TaskQueue.from_list([
+            {"query": "Task 1", "environment_data": {}},
+            {"query": "Task 2", "environment_data": {}},
+        ])
+        benchmark = DummyBenchmark()
+        benchmark.run(tasks, agent_data={"model": "test"})
+
+        # Both tasks should have produced reports with usage
+        assert len(benchmark.reports) == 2
+        assert "usage" in benchmark.reports[0]
+        assert "usage" in benchmark.reports[1]
+
+    def test_usage_property_returns_total(self):
+        """benchmark.usage returns the running total."""
+        from conftest import DummyBenchmark
+
+        tasks = TaskQueue.from_list([{"query": "Test", "environment_data": {}}])
+        benchmark = DummyBenchmark()
+        benchmark.run(tasks, agent_data={"model": "test"})
+
+        # usage property should return a Usage object (even if empty)
+        total = benchmark.usage
+        assert total is not None
+        # cost may be None if DummyModelAdapter doesn't provide usage

tests/test_core/test_registry.py

@@ -256,3 +256,230 @@ def worker(worker_id: int):
         for worker_id, traces in results.items():
             assert f"agent_{worker_id}" in traces["agents"]
             assert len(traces["agents"]) == 1
+
+
+# ==================== Usage Tracking Tests ====================
+
+
+class MockUsageComponent(TraceableMixin):
+    """Component that implements UsageTrackableMixin for testing."""
+
+    def __init__(self, name: str, cost: float = 0.0, input_tokens: int = 0, output_tokens: int = 0):
+        super().__init__()
+        self._name = name
+        self._cost = cost
+        self._input_tokens = input_tokens
+        self._output_tokens = output_tokens
+
+    def gather_traces(self) -> Dict[str, Any]:
+        return {"name": self._name}
+
+    def gather_usage(self):
+        from maseval.core.usage import TokenUsage
+        return TokenUsage(
+            cost=self._cost,
+            input_tokens=self._input_tokens,
+            output_tokens=self._output_tokens,
+            total_tokens=self._input_tokens + self._output_tokens,
+        )
+
+
+class MockBrokenUsageComponent(TraceableMixin):
+    """Component whose gather_usage raises an exception."""
+
+    def __init__(self):
+        super().__init__()
+
+    def gather_traces(self) -> Dict[str, Any]:
+        return {}
+
+    def gather_usage(self):
+        raise RuntimeError("Usage collection failed")
+
+
+# Ensure MockUsageComponent also inherits UsageTrackableMixin
+from maseval.core.usage import UsageTrackableMixin
+
+
+class UsageAwareComponent(TraceableMixin, UsageTrackableMixin):
+    """Component with both tracing and usage tracking."""
+
+    def __init__(self, cost: float = 0.0, input_tokens: int = 0, output_tokens: int = 0):
+        TraceableMixin.__init__(self)
+        self._cost = cost
+        self._input_tokens = input_tokens
+        self._output_tokens = output_tokens
+
+    def gather_traces(self) -> Dict[str, Any]:
+        return {"traced": True}
+
+    def gather_usage(self):
+        from maseval.core.usage import TokenUsage
+        return TokenUsage(
+            cost=self._cost,
+            input_tokens=self._input_tokens,
+            output_tokens=self._output_tokens,
+            total_tokens=self._input_tokens + self._output_tokens,
+        )
+
+
+class BrokenUsageComponent(TraceableMixin, UsageTrackableMixin):
+    """Component whose gather_usage raises an exception."""
+
+    def __init__(self):
+        TraceableMixin.__init__(self)
+
+    def gather_traces(self) -> Dict[str, Any]:
+        return {}
+
+    def gather_usage(self):
+        raise RuntimeError("Usage collection failed")
+
+
+@pytest.mark.core
+class TestRegistryUsageCollection:
+    """Tests for usage tracking through the component registry."""
+
+    def test_register_usage_trackable_component(self):
+        """UsageTrackableMixin component is registered in the usage registry."""
+        registry = ComponentRegistry()
+        component = UsageAwareComponent(cost=0.05, input_tokens=100, output_tokens=50)
+
+        registry.register("models", "main_model", component)
+
+        assert "models:main_model" in registry._usage_registry
+        assert registry._usage_registry["models:main_model"] is component
+
+    def test_non_usage_component_not_in_usage_registry(self):
+        """Components without UsageTrackableMixin are NOT in the usage registry."""
+        registry = ComponentRegistry()
+        component = MockTraceableComponent("test")
+
+        registry.register("agents", "my_agent", component)
+
+        assert "agents:my_agent" in registry._trace_registry
+        assert "agents:my_agent" not in registry._usage_registry
+
+    def test_collect_usage_basic(self):
+        """collect_usage returns structured dict with usage from registered components."""
+        from maseval.core.usage import TokenUsage
+
+        registry = ComponentRegistry()
+        model = UsageAwareComponent(cost=0.10, input_tokens=500, output_tokens=200)
+        registry.register("models", "main_model", model)
+
+        usage = registry.collect_usage()
+
+        assert "metadata" in usage
+        assert "models" in usage
+        assert "main_model" in usage["models"]
+
+        model_usage = usage["models"]["main_model"]
+        assert model_usage["cost"] == 0.10
+        assert model_usage["input_tokens"] == 500
+        assert model_usage["output_tokens"] == 200
+        assert model_usage["total_tokens"] == 700
+
+    def test_collect_usage_multiple_components(self):
+        """Multiple components across categories are all collected."""
+        registry = ComponentRegistry()
+        model = UsageAwareComponent(cost=0.10, input_tokens=500, output_tokens=200)
+        tool = UsageAwareComponent(cost=0.05, input_tokens=0, output_tokens=0)
+
+        registry.register("models", "main_model", model)
+        registry.register("tools", "search_tool", tool)
+
+        usage = registry.collect_usage()
+
+        assert "main_model" in usage["models"]
+        assert "search_tool" in usage["tools"]
+        assert usage["models"]["main_model"]["cost"] == 0.10
+        assert usage["tools"]["search_tool"]["cost"] == 0.05
+
+    def test_collect_usage_injects_grouping_fields(self):
+        """Registry injects category and component_name into usage records."""
+        registry = ComponentRegistry()
+        model = UsageAwareComponent(cost=0.10, input_tokens=100, output_tokens=50)
+        registry.register("models", "main_model", model)
+
+        usage = registry.collect_usage()
+
+        model_usage = usage["models"]["main_model"]
+        assert model_usage["category"] == "models"
+        assert model_usage["component_name"] == "main_model"
+
+    def test_total_usage_accumulates(self):
+        """total_usage property reflects accumulated usage across collect_usage calls."""
+        registry = ComponentRegistry()
+        model = UsageAwareComponent(cost=0.10, input_tokens=100, output_tokens=50)
+        registry.register("models", "main_model", model)
+
+        # First collection
+        registry.collect_usage()
+        total1 = registry.total_usage
+        assert total1.cost == pytest.approx(0.10)
+
+        # Clear and re-register (simulates next repetition)
+        registry.clear()
+        model2 = UsageAwareComponent(cost=0.20, input_tokens=200, output_tokens=100)
+        registry.register("models", "main_model", model2)
+
+        # Second collection
+        registry.collect_usage()
+        total2 = registry.total_usage
+        assert total2.cost == pytest.approx(0.30)
+
+    def test_usage_by_component_accumulates(self):
+        """usage_by_component accumulates per key across repetitions."""
+        registry = ComponentRegistry()
+        model = UsageAwareComponent(cost=0.10, input_tokens=100, output_tokens=50)
+        registry.register("models", "main_model", model)
+        registry.collect_usage()
+
+        # Clear and re-register for second repetition
+        registry.clear()
+        model2 = UsageAwareComponent(cost=0.20, input_tokens=200, output_tokens=100)
+        registry.register("models", "main_model", model2)
+        registry.collect_usage()
+
+        by_comp = registry.usage_by_component
+        assert "models:main_model" in by_comp
+
+        total = by_comp["models:main_model"]
+        assert total.input_tokens == 300
+        assert total.output_tokens == 150
+        assert total.cost == pytest.approx(0.30)
+
+    def test_usage_persists_across_clear(self):
+        """clear() does NOT reset total_usage or usage_by_component."""
+        registry = ComponentRegistry()
+        model = UsageAwareComponent(cost=0.10, input_tokens=100, output_tokens=50)
+        registry.register("models", "main_model", model)
+        registry.collect_usage()
+
+        # Clear only removes per-repetition state
+        registry.clear()
+
+        assert registry.total_usage.cost == pytest.approx(0.10)
+        assert "models:main_model" in registry.usage_by_component
+
+    def test_collect_usage_handles_error_gracefully(self):
+        """If gather_usage raises, the error is captured in the usage dict."""
+        registry = ComponentRegistry()
+        broken = BrokenUsageComponent()
+        registry.register("models", "bad_model", broken)
+
+        usage = registry.collect_usage()
+
+        assert "bad_model" in usage["models"]
+        assert "error" in usage["models"]["bad_model"]
+        assert "RuntimeError" in usage["models"]["bad_model"]["error_type"]
+
+    def test_collect_usage_empty_registry(self):
+        """collect_usage with no components returns empty structure."""
+        registry = ComponentRegistry()
+        usage = registry.collect_usage()
+
+        assert usage["metadata"]["total_components"] == 0
+        assert usage["models"] == {}
+        assert usage["agents"] == {}