from botocore.exceptions import CapacityNotAvailableError
from botocore.retries import bucket
from tests import unittest


class FakeClock(bucket.Clock):
    def __init__(self, timestamp_sequences):
        self.timestamp_sequences = timestamp_sequences
        self.sleep_call_amounts = []

    def sleep(self, amount):
        self.sleep_call_amounts.append(amount)

    def current_time(self):
        return self.timestamp_sequences.pop(0)


class TestTokenBucket(unittest.TestCase):
    def setUp(self):
        self.timestamp_sequences = [0]
        self.clock = FakeClock(self.timestamp_sequences)

    def create_token_bucket(self, max_rate=10, min_rate=0.1):
        return bucket.TokenBucket(
            max_rate=max_rate, clock=self.clock, min_rate=min_rate
        )

    def test_can_acquire_amount(self):
        self.timestamp_sequences.extend(
            [
                # Requests tokens every second, which is well below our
                # 10 TPS fill rate.
                1,
                2,
                3,
                4,
                5,
            ]
        )
        token_bucket = self.create_token_bucket(max_rate=10)
        for _ in range(5):
            self.assertTrue(token_bucket.acquire(1, block=False))

    def test_can_change_max_capacity_lower(self):
        # Requests at 1 TPS.
        self.timestamp_sequences.extend([1, 2, 3, 4, 5])
        token_bucket = self.create_token_bucket(max_rate=10)
        # Request the first 5 tokens with max_rate=10
        for _ in range(5):
            self.assertTrue(token_bucket.acquire(1, block=False))
        # Now scale the max_rate down to 1 on the 5th second.
        self.timestamp_sequences.append(5)
        token_bucket.max_rate = 1
        # And then from seconds 6-10 we request at one per second.
        self.timestamp_sequences.extend([6, 7, 8, 9, 10])
        for _ in range(5):
            self.assertTrue(token_bucket.acquire(1, block=False))

    def test_max_capacity_is_at_least_one(self):
        token_bucket = self.create_token_bucket()
        self.timestamp_sequences.append(1)
        token_bucket.max_rate = 0.5
        self.assertEqual(token_bucket.max_rate, 0.5)
        self.assertEqual(token_bucket.max_capacity, 1)

    def test_acquire_fails_on_non_block_mode_returns_false(self):
        self.timestamp_sequences.extend(
            [
                # Initial creation time.
                0,
                # Requests a token 1 second later.
                1,
            ]
        )
        token_bucket = self.create_token_bucket(max_rate=10)
        with self.assertRaises(CapacityNotAvailableError):
            token_bucket.acquire(100, block=False)

    def test_can_retrieve_at_max_send_rate(self):
        self.timestamp_sequences.extend(
            [
                # Request a new token every 100ms (10 TPS) for 2 seconds.
                1 + 0.1 * i
                for i in range(20)
            ]
        )
        token_bucket = self.create_token_bucket(max_rate=10)
        for _ in range(20):
            self.assertTrue(token_bucket.acquire(1, block=False))

    def test_acquiring_blocks_when_capacity_reached(self):
        # This is 1 token every 0.1 seconds.
        token_bucket = self.create_token_bucket(max_rate=10)
        self.timestamp_sequences.extend(
            [
                # The first acquire() happens after .1 seconds.
                0.1,
                # The second acquire() will fail because we get tokens at
                # 1 per 0.1 seconds.  We will then sleep for 0.05 seconds until we
                # get a new token.
                0.15,
                # And at 0.2 seconds we get our token.
                0.2,
                # And at 0.3 seconds we have no issues getting a token.
                # Because we're using such small units (to avoid bloating the
                # test run time), we have to go slightly over 0.3 seconds here.
                0.300001,
            ]
        )
        self.assertTrue(token_bucket.acquire(1, block=False))
        self.assertEqual(token_bucket.available_capacity, 0)
        self.assertTrue(token_bucket.acquire(1, block=True))
        self.assertEqual(token_bucket.available_capacity, 0)
        self.assertTrue(token_bucket.acquire(1, block=False))

    def test_rate_cant_go_below_min(self):
        token_bucket = self.create_token_bucket(max_rate=1, min_rate=0.2)
        self.timestamp_sequences.append(1)
        token_bucket.max_rate = 0.1
        self.assertEqual(token_bucket.max_rate, 0.2)
        self.assertEqual(token_bucket.max_capacity, 1)