from energiamesh.core import SafetyBudget, PrivacyBudget


def test_safety_budget_basic():
    sb = SafetyBudget(enabled=True, max_current_draw_a=50.0, max_voltage_variation_pu=0.02)
    assert sb.enabled is True
    assert sb.max_current_draw_a == 50.0
    assert sb.max_voltage_variation_pu == 0.02
    # Update device limits
    sb.update("DER-01", 40.0)
    assert sb.device_limits["DER-01"] == 40.0


def test_privacy_budget_basic_and_use():
    pb = PrivacyBudget(enabled=True, allowed_signals=["forecast"], total_budget_units=1.0)
    assert pb.enabled is True
    assert pb.allowed_signals == ["forecast"]
    pb.per_signal_budget["forecast"] = 1.0
    pb.use("forecast", 0.25)
    assert pb.per_signal_budget["forecast"] == 0.75