import time
import unittest

from danswer.connectors.cross_connector_utils.rate_limit_wrapper import (
    rate_limit_builder,
)


class TestRateLimit(unittest.TestCase):
    call_cnt = 0

    def test_rate_limit_basic(self) -> None:
        self.call_cnt = 0

        @rate_limit_builder(max_calls=2, period=5)
        def func() -> None:
            self.call_cnt += 1

        start = time.time()

        # make calls that shouldn't be rate-limited
        func()
        func()
        time_to_finish_non_ratelimited = time.time() - start

        # make a call which SHOULD be rate-limited
        func()
        time_to_finish_ratelimited = time.time() - start

        self.assertEqual(self.call_cnt, 3)
        self.assertLess(time_to_finish_non_ratelimited, 1)
        self.assertGreater(time_to_finish_ratelimited, 5)


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