from __future__ import annotations

import unittest

from datahub.governance.circuit import CircuitBreaker
from datahub.governance.ratelimit import TokenBucket
from datahub.governance.retry import RetryError, retry_call


class FakeClock:
    def __init__(self) -> None:
        self.value = 0.0

    def __call__(self) -> float:
        return self.value


class GovernanceTests(unittest.TestCase):
    def test_token_bucket_caps_burst_at_capacity(self) -> None:
        clock = FakeClock()
        bucket = TokenBucket(rate_per_minute=300, capacity=300, clock=clock)
        ok = 0
        for _ in range(400):
            if bucket.acquire(block=False):
                ok += 1
        self.assertEqual(ok, 300)
        clock.value = 60
        self.assertTrue(bucket.acquire(block=False))

    def test_circuit_opens_after_five_failures_and_half_opens(self) -> None:
        clock = FakeClock()
        breaker = CircuitBreaker(clock=clock, open_seconds=120)
        for _ in range(5):
            breaker.record_failure("boom")
        self.assertEqual(breaker.snapshot().state, "open")
        self.assertFalse(breaker.allow())
        clock.value = 120
        self.assertEqual(breaker.snapshot().state, "half_open")
        self.assertTrue(breaker.allow())
        breaker.record_success()
        self.assertEqual(breaker.snapshot().state, "closed")

    def test_retry_exhausts(self) -> None:
        calls = {"n": 0}

        def fail():
            calls["n"] += 1
            raise RuntimeError("no")

        with self.assertRaises(RetryError):
            retry_call(fail, attempts=3, sleeper=lambda _d: None)
        self.assertEqual(calls["n"], 3)


if __name__ == "__main__":
    unittest.main()