import unittest
from guardrail_space.contract import SafetyContract
from guardrail_space.policy import PolicyEngine


class TestPolicyEngine(unittest.TestCase):
    def test_budget_and_pref(self):
        c = SafetyContract(
            contract_id="policy-001",
            pre_conditions=["state['speed'] <= 2.0"],
            budgets={"time": 5.0, "energy": 50.0},
            collision_rules=["state['distance_to_obstacle'] >= 1.0"],
        )
        engine = PolicyEngine(c)
        plan = {"action": "move", "costs": {"time": 3.0, "energy": 20.0}, "speed": 1.0}
        self.assertTrue(engine.evaluate(plan, {"speed": 1.0, "distance_to_obstacle": 2}).get("allowed"))
        plan2 = {"action": "move", "costs": {"time": 6.0, "energy": 20.0}, "speed": 1.0}
        self.assertFalse(engine.evaluate(plan2, {"speed": 1.0, "distance_to_obstacle": 2}).get("allowed"))