# ~ encoding: utf-8  ~
import time
import unittest
from mock import MagicMock, patch


import tbbackoff as tbb


class RateLimitTestCase(unittest.TestCase):

    def setUp(self):

        # Stop the time
        self.time_patcher = patch('tbbackoff.time')
        self.time_mock = self.time_patcher.start()
        self.time_mock.time.return_value = time.time()

        def sleep_mock(i):
            self.time_mock.time.return_value += i

        self.time_mock.sleep.side_effect = sleep_mock

    def tearDown(self):
        self.time_patcher.stop()

    def test_first_request_is_not_ratelimited(self):
        def ratelimited_function():
            pass

        ratelimit = tbb.RateLimit()
        function = ratelimit(ratelimited_function)

        function()

        self.assertFalse(self.time_mock.sleep.called)

    def test_ratelimit_after_bucket_empty(self):
        def ratelimited_function():
            pass

        ratelimit = tbb.RateLimit(5, token_time=0.5)
        function = ratelimit(ratelimited_function)

        # Se ejecuta hasta el borde del ratelimit
        for i in range(0, ratelimit.bucket_fill):
            function()

        # No hubo ratelimit
        self.assertFalse(self.time_mock.sleep.called)

        # Al ejecutar mas, tiene que hacer ratelimit
        function()
        self.assertTrue(self.time_mock.sleep.called)

        # Si pasa el tiempo debería llenarse el bucket
        self.time_mock.time.return_value += 3
        function()

        # El bucket está lleno - 1 token
        self.assertEqual(ratelimit.bucket_fill, 4)

    def test_bucket_is_emptied_on_backpressure(self):
        ''' Si se obtiene un 429 el bucket size se reduce
        '''

        empty_bucket_function = MagicMock()

        ratelimit = tbb.RateLimit(6, token_time=0.5, empty_bucket_callback=empty_bucket_function)

        function_mock = MagicMock()
        function_mock.side_effect = [tbb.RateLimitBackPressure(), 'OK']
        function_mock.__name__ = 'function'
        function = ratelimit(function_mock)

        function()

        # El bucket se vacía
        self.assertTrue(empty_bucket_function.called)

    def test_backoff_on_backpressure(self):

        backoff_callback_fn = MagicMock()

        ratelimit = tbb.RateLimit(6, token_time=1, backoff_callback=backoff_callback_fn)

        pressure_response = [tbb.RateLimitBackPressure()]
        ok_response = ['OK']

        function_mock = MagicMock()
        function_mock.side_effect = pressure_response * 3 + ok_response * 8
        function_mock.__name__ = 'function'
        function = ratelimit(function_mock)

        # Al llamar a la función causa 3 backpressure, se recibe la respuesta
        response = function()
        self.assertEqual(response, 'OK')

        # Y se incrementó el backoff
        self.assertEqual(ratelimit.backoff, 4)
        self.assertEqual(backoff_callback_fn.call_count, 3)

        # A medida que pasa el tiempo con llamadas sin backpressure, el backoff se va reduciendo
        self.time_mock.time.return_value += ratelimit.backoff * 8
        high_backoff = ratelimit.backoff

        for i in range(4):
            function()

        self.assertLess(ratelimit.backoff, high_backoff)

        # El backoff se ejecutó cuatro veces más para achicarlo
        self.assertEqual(backoff_callback_fn.call_count, 7)

        # Hasta que llega a cero
        function()
        self.assertEqual(ratelimit.backoff, 0)

        # Y no se reduce mas
        function()
        self.assertEqual(ratelimit.backoff, 0)

    def test_max_retries(self):
        def bad_function():
            raise tbb.RateLimitBackPressure()

        function = tbb.RateLimit(max_retries=4)(bad_function)

        with self.assertRaises(tbb.MaxRetriesReached):
            function()
