import threading
import time
import unittest

from singleflight import SingleFlight


class TestSingleFlight(unittest.TestCase):
    def test_single_call_runs_fn(self):
        sf = SingleFlight()
        calls = {"n": 0}

        def fn():
            calls["n"] += 1
            return 42

        self.assertEqual(sf.do("k", fn), 42)
        self.assertEqual(calls["n"], 1)

    def test_concurrent_coalesces(self):
        sf = SingleFlight()
        calls = {"n": 0}
        lock = threading.Lock()
        barrier = threading.Barrier(8)

        def fn():
            with lock:
                calls["n"] += 1
            time.sleep(0.2)
            return "result"

        results = [None] * 8
        errors = []

        def worker(i):
            barrier.wait()
            try:
                results[i] = sf.do("k", fn)
            except BaseException as e:  # noqa
                errors.append(e)

        ts = [threading.Thread(target=worker, args=(i,)) for i in range(8)]
        for t in ts:
            t.start()
        for t in ts:
            t.join()
        self.assertFalse(errors)
        self.assertTrue(all(r == "result" for r in results))
        self.assertEqual(calls["n"], 1)

    def test_exception_propagates_to_all(self):
        sf = SingleFlight()
        calls = {"n": 0}

        def fn():
            calls["n"] += 1
            time.sleep(0.2)
            raise ValueError("boom")

        results = []
        errors = []
        barrier = threading.Barrier(5)

        def worker():
            barrier.wait()
            try:
                sf.do("k", fn)
            except Exception as e:  # noqa
                errors.append(e)
            else:
                results.append("ok")

        ts = [threading.Thread(target=worker) for _ in range(5)]
        for t in ts:
            t.start()
        for t in ts:
            t.join()
        self.assertEqual(len(errors), 5)
        self.assertTrue(all(isinstance(e, ValueError) and str(e) == "boom" for e in errors))
        self.assertEqual(calls["n"], 1)
        # Same exception object shared.
        self.assertTrue(len({id(e) for e in errors}) == 1)

    def test_subsequent_call_runs_again(self):
        sf = SingleFlight()
        calls = {"n": 0}

        def fn():
            calls["n"] += 1
            return calls["n"]

        self.assertEqual(sf.do("k", fn), 1)
        self.assertEqual(sf.do("k", fn), 2)
        self.assertEqual(calls["n"], 2)

    def test_different_keys_independent(self):
        sf = SingleFlight()
        calls = {"a": 0, "b": 0}

        def make(key):
            def fn():
                calls[key] += 1
                time.sleep(0.1)
                return key
            return fn

        r = {}
        def wa():
            r["a"] = sf.do("a", make("a"))
        def wb():
            r["b"] = sf.do("b", make("b"))
        ta = threading.Thread(target=wa); tb = threading.Thread(target=wb)
        ta.start(); tb.start(); ta.join(); tb.join()
        self.assertEqual(r["a"], "a")
        self.assertEqual(r["b"], "b")
        self.assertEqual(calls["a"], 1)
        self.assertEqual(calls["b"], 1)

    def test_inflight_count(self):
        sf = SingleFlight()
        started = threading.Event()
        release = threading.Event()

        def fn():
            started.set()
            release.wait()
            return 1

        def leader():
            sf.do("k", fn)

        t = threading.Thread(target=leader)
        t.start()
        started.wait()
        self.assertEqual(sf.inflight_count, 1)
        release.set()
        t.join()
        self.assertEqual(sf.inflight_count, 0)

    def test_leader_crash_does_not_leak_key(self):
        sf = SingleFlight()

        def fn():
            raise RuntimeError("crash")

        try:
            sf.do("k", fn)
        except RuntimeError:
            pass
        # Key freed: next call runs fn again.
        calls = {"n": 0}

        def fn2():
            calls["n"] += 1
            return "ok"

        self.assertEqual(sf.do("k", fn2), "ok")
        self.assertEqual(calls["n"], 1)
        self.assertEqual(sf.inflight_count, 0)


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