from __future__ import annotations
|
|
from pathlib import Path
|
from datetime import datetime, timezone
|
import tempfile
|
import unittest
|
|
from hibor_fast_collection.budget import Budget, BudgetExpired
|
from hibor_fast_collection.models import (
|
ContractError, ErrorCode, EvidenceState, PackageState, TaskSpec, TerminalReceipt,
|
)
|
from helpers import task_spec
|
|
|
class FakeClock:
|
def __init__(self): self.value = 0.0
|
def __call__(self): return self.value
|
|
|
class ModelsBudgetTests(unittest.TestCase):
|
def test_task_spec_exact_and_batch_budget(self):
|
with tempfile.TemporaryDirectory() as tmp:
|
spec = TaskSpec.from_mapping(task_spec(Path(tmp), quantity=3))
|
self.assertEqual(spec.total_budget_ms, 1_080_000)
|
bad = task_spec(Path(tmp))
|
bad["package_name"] = "other"
|
with self.assertRaises(ContractError): TaskSpec.from_mapping(bad)
|
|
def test_key_order_rejected(self):
|
with tempfile.TemporaryDirectory() as tmp:
|
value = task_spec(Path(tmp))
|
first = value.pop("schema_version")
|
value["schema_version"] = first
|
with self.assertRaises(ContractError): TaskSpec.from_mapping(value)
|
|
def test_budget_work_and_close_deadline(self):
|
clock = FakeClock()
|
budget = Budget(10_000, 2_000, clock=clock)
|
self.assertEqual(budget.require("x", "y", 9_000), 8_000)
|
clock.value = 8.001
|
with self.assertRaises(BudgetExpired): budget.require("x", "y", 1_000)
|
self.assertGreater(budget.require("close", "z", 1_000, close=True), 0)
|
clock.value = 10.1
|
with self.assertRaises(BudgetExpired): budget.require("close", "z", 1_000, close=True)
|
|
def test_observed_deadline_and_batch_increment(self):
|
fresh_clock = FakeClock()
|
fresh = Budget(840_000, 30_000, batch_increment_ms=240_000,
|
clock=fresh_clock)
|
fresh.begin_item(1)
|
self.assertEqual((fresh.remaining_ms(), fresh.remaining_ms(close=True),
|
fresh.remaining_ms(batch_close=True)),
|
(570_000, 600_000, 840_000))
|
|
clock = FakeClock()
|
wall = lambda: datetime(2026, 7, 29, 0, 5, tzinfo=timezone.utc)
|
budget = Budget(840_000, 30_000, observed_at_utc="2026-07-29T00:00:00Z",
|
batch_increment_ms=240_000, clock=clock, utc_now=wall)
|
budget.begin_item(1)
|
self.assertEqual(budget.remaining_ms(), 270_000)
|
self.assertEqual(budget.remaining_ms(close=True), 300_000)
|
self.assertEqual(budget.remaining_ms(batch_close=True), 540_000)
|
with self.assertRaises(ValueError):
|
budget.begin_item(2)
|
clock.value = 299.0
|
self.assertTrue(budget.complete_item(1))
|
budget.begin_item(2)
|
self.assertEqual(budget.remaining_ms(), 210_000)
|
self.assertEqual(budget.remaining_ms(close=True), 240_000)
|
clock.value = 509.1
|
with self.assertRaises(BudgetExpired): budget.checkpoint("item_2", "terminal")
|
self.assertGreater(budget.remaining_ms(close=True), 0)
|
clock.value = 539.1
|
self.assertFalse(budget.complete_item(2))
|
|
def test_terminal_receipt_cross_fields(self):
|
target = Path("terminal.json")
|
valid = TerminalReceipt(target, target, True, True, EvidenceState.V, True, False, True,
|
10, "a"*64, None, PackageState.P09_TIMING)
|
valid.validate()
|
file_u = TerminalReceipt(target, None, True, None, EvidenceState.U, True, False, True,
|
None, None, ErrorCode.RECOVERY_UNKNOWN, PackageState.P09_TIMING)
|
file_u.validate()
|
live_u = TerminalReceipt(target, None, True, None, EvidenceState.U, None, True, True,
|
None, None, ErrorCode.PROCESS_LIVENESS_UNKNOWN, PackageState.P09_TIMING)
|
live_u.validate()
|
with self.assertRaises(ContractError):
|
TerminalReceipt(target, None, False, False, EvidenceState.U, False, False, True,
|
None, None, ErrorCode.RECOVERY_UNKNOWN, PackageState.P09_TIMING).validate()
|
|
|
if __name__ == "__main__": unittest.main()
|
