"""Tests for batch-level energy accounting.""" from __future__ import annotations import re from contextlib import contextmanager from dataclasses import dataclass from typing import Generator import pytest from openjarvis.telemetry.batch import BatchMetrics, EnergyBatch # --------------------------------------------------------------------------- # BatchMetrics defaults # --------------------------------------------------------------------------- class TestBatchMetricsDefaults: def test_all_defaults(self) -> None: m = BatchMetrics() assert m.batch_id == "" assert m.total_requests == 0 assert m.total_tokens == 0 assert m.total_energy_joules == 0.0 assert m.energy_per_token_joules == 0.0 assert m.energy_per_request_joules == 0.0 assert m.mean_power_watts == 0.0 assert m.mean_throughput_tok_per_sec == 0.0 assert m.per_request_energy == [] def test_custom_values(self) -> None: m = BatchMetrics( batch_id="abc", total_requests=5, total_tokens=100, total_energy_joules=10.0, energy_per_token_joules=0.1, energy_per_request_joules=2.0, mean_power_watts=50.0, mean_throughput_tok_per_sec=200.0, per_request_energy=[1.0, 2.0, 3.0, 2.5, 1.5], ) assert m.batch_id == "abc" assert m.total_requests == 5 assert m.per_request_energy == [1.0, 2.0, 3.0, 2.5, 1.5] # --------------------------------------------------------------------------- # Batch ID generation # --------------------------------------------------------------------------- class TestBatchIdGeneration: def test_auto_generated_uuid(self) -> None: batch = EnergyBatch() # UUID4 format: 8-4-4-4-12 hex digits uuid4_re = ( r"^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}" r"-[89ab][0-9a-f]{3}-[0-9a-f]{12}$" ) assert re.match(uuid4_re, batch.batch_id) def test_custom_batch_id(self) -> None: batch = EnergyBatch(batch_id="my-batch-42") assert batch.batch_id == "my-batch-42" def test_unique_ids(self) -> None: ids = {EnergyBatch().batch_id for _ in range(100)} assert len(ids) == 100 # --------------------------------------------------------------------------- # EnergyBatch without monitor # --------------------------------------------------------------------------- class TestEnergyBatchNoMonitor: def test_record_request_accumulation(self) -> None: batch = EnergyBatch() with batch.sample() as ctx: ctx.record_request(tokens=50) ctx.record_request(tokens=30) ctx.record_request(tokens=20) assert batch.metrics is not None assert batch.metrics.total_requests == 3 assert batch.metrics.total_tokens == 100 def test_energy_stays_zero_without_monitor(self) -> None: batch = EnergyBatch() with batch.sample() as ctx: ctx.record_request(tokens=50) assert batch.metrics is not None assert batch.metrics.total_energy_joules == 0.0 assert batch.metrics.energy_per_token_joules == 0.0 assert batch.metrics.mean_power_watts == 0.0 def test_per_request_energy_from_record(self) -> None: """When no monitor, per-request energy comes from record_request calls.""" batch = EnergyBatch() with batch.sample() as ctx: ctx.record_request(tokens=50, energy_joules=1.0) ctx.record_request(tokens=30, energy_joules=2.0) assert batch.metrics is not None assert batch.metrics.per_request_energy == [1.0, 2.0] assert batch.metrics.total_energy_joules == 3.0 def test_metrics_computed_on_exit(self) -> None: batch = EnergyBatch() assert batch.metrics is None # Before sample() with batch.sample() as ctx: ctx.record_request(tokens=100) assert batch.metrics is not None def test_no_requests_yields_zero_metrics(self) -> None: batch = EnergyBatch() with batch.sample() as _ctx: pass # No requests recorded m = batch.metrics assert m is not None assert m.total_requests == 0 assert m.total_tokens == 0 assert m.energy_per_token_joules == 0.0 assert m.energy_per_request_joules == 0.0 def test_throughput_computed(self) -> None: batch = EnergyBatch() with batch.sample() as ctx: ctx.record_request(tokens=1000) assert batch.metrics is not None assert batch.metrics.mean_throughput_tok_per_sec > 0 # --------------------------------------------------------------------------- # EnergyBatch with mock monitor # --------------------------------------------------------------------------- @dataclass class _FakeEnergySample: energy_joules: float = 0.0 mean_power_watts: float = 0.0 class _FakeMonitor: """Minimal mock that mimics EnergyMonitor.sample() context manager.""" def __init__(self, energy_joules: float = 10.0, mean_power_watts: float = 200.0): self._energy = energy_joules self._power = mean_power_watts @contextmanager def sample(self) -> Generator[_FakeEnergySample, None, None]: s = _FakeEnergySample() yield s s.energy_joules = self._energy s.mean_power_watts = self._power class TestEnergyBatchWithMonitor: def test_energy_from_monitor(self) -> None: monitor = _FakeMonitor(energy_joules=10.0, mean_power_watts=200.0) batch = EnergyBatch(energy_monitor=monitor) with batch.sample() as ctx: ctx.record_request(tokens=100) m = batch.metrics assert m is not None assert m.total_energy_joules == pytest.approx(10.0) assert m.mean_power_watts == pytest.approx(200.0) def test_energy_per_token_with_monitor(self) -> None: monitor = _FakeMonitor(energy_joules=20.0) batch = EnergyBatch(energy_monitor=monitor) with batch.sample() as ctx: ctx.record_request(tokens=100) ctx.record_request(tokens=100) m = batch.metrics assert m is not None assert m.total_tokens == 200 assert m.energy_per_token_joules == pytest.approx(20.0 / 200.0) def test_energy_per_request_with_monitor(self) -> None: monitor = _FakeMonitor(energy_joules=15.0) batch = EnergyBatch(energy_monitor=monitor) with batch.sample() as ctx: ctx.record_request(tokens=50) ctx.record_request(tokens=50) ctx.record_request(tokens=50) m = batch.metrics assert m is not None assert m.total_requests == 3 assert m.energy_per_request_joules == pytest.approx(15.0 / 3.0) def test_batch_id_in_metrics(self) -> None: batch = EnergyBatch(batch_id="test-batch-99") with batch.sample() as ctx: ctx.record_request(tokens=10) assert batch.metrics is not None assert batch.metrics.batch_id == "test-batch-99" # --------------------------------------------------------------------------- # Energy per token calculation # --------------------------------------------------------------------------- class TestEnergyPerToken: def test_basic_division(self) -> None: monitor = _FakeMonitor(energy_joules=50.0) batch = EnergyBatch(energy_monitor=monitor) with batch.sample() as ctx: ctx.record_request(tokens=500) assert batch.metrics is not None assert batch.metrics.energy_per_token_joules == pytest.approx(50.0 / 500.0) def test_zero_tokens_yields_zero(self) -> None: monitor = _FakeMonitor(energy_joules=10.0) batch = EnergyBatch(energy_monitor=monitor) with batch.sample() as _ctx: pass # No requests assert batch.metrics is not None assert batch.metrics.energy_per_token_joules == 0.0 def test_many_small_requests(self) -> None: monitor = _FakeMonitor(energy_joules=1.0) batch = EnergyBatch(energy_monitor=monitor) with batch.sample() as ctx: for _ in range(100): ctx.record_request(tokens=10) m = batch.metrics assert m is not None assert m.total_tokens == 1000 assert m.total_requests == 100 assert m.energy_per_token_joules == pytest.approx(1.0 / 1000.0) # --------------------------------------------------------------------------- # Per-request energy list tracking # --------------------------------------------------------------------------- class TestPerRequestEnergy: def test_tracks_each_request(self) -> None: batch = EnergyBatch() with batch.sample() as ctx: ctx.record_request(tokens=10, energy_joules=0.5) ctx.record_request(tokens=20, energy_joules=1.0) ctx.record_request(tokens=30, energy_joules=1.5) m = batch.metrics assert m is not None assert m.per_request_energy == [0.5, 1.0, 1.5] assert len(m.per_request_energy) == 3 def test_empty_when_no_requests(self) -> None: batch = EnergyBatch() with batch.sample() as _ctx: pass assert batch.metrics is not None assert batch.metrics.per_request_energy == [] def test_zeros_when_no_per_request_energy(self) -> None: batch = EnergyBatch() with batch.sample() as ctx: ctx.record_request(tokens=10) ctx.record_request(tokens=20) m = batch.metrics assert m is not None assert m.per_request_energy == [0.0, 0.0]