<commit_before>import unittest

import requests

from garage.http2 import clients
from garage.http2 import policies


class TestClient(unittest.TestCase):

    def test_client(self):
        req_to_rep = {
            ('GET', 'uri_1'): (200, 'hello world'),
            ('GET', 'uri_2'): Exception('error_2'),
            ('GET', 'uri_3'): (200, 'GET uri_3'),
            ('HEAD', 'uri_3'): (200, 'HEAD uri_3'),
            ('POST', 'uri_3'): (200, 'POST uri_3'),
            ('PUT', 'uri_3'): (200, 'PUT uri_3'),
            ('GET', 'uri_4'): (404, ''),
        }
        client = clients.Client(
            _session=MockSession(req_to_rep),
            _sleep=fake_sleep,
        )

        rep = client.send(clients.Request('GET', 'uri_1'))
        self.assertEqual('hello world', rep.content)

        with self.assertRaisesRegex(Exception, 'error_2'):
            client.send(clients.Request('GET', 'uri_2'))

        self.assertEqual('GET uri_3', client.get('uri_3').content)
        self.assertEqual('HEAD uri_3', client.head('uri_3').content)
        self.assertEqual('POST uri_3', client.post('uri_3').content)
        self.assertEqual('PUT uri_3', client.put('uri_3').content)

        with self.assertRaisesRegex(requests.HTTPError, 'http error'):
            client.get('uri_4')

    def test_no_retry(self):
        session = MockSession({('GET', 'uri_1'): (400, 'error!')})
        client = clients.Client(
            retry_policy=policies.NoRetry(),
            _session=session,
            _sleep=fake_sleep,
        )
        with self.assertRaisesRegex(requests.HTTPError, 'http error'):
            client.get('uri_1')
        self.assertEqual(1, len(session._logs))
        for req in session._logs:
            self.assertTrue(isinstance(req, requests.Request))
            self.assertEqual('GET', req.method)
            self.assertEqual('uri_1', req.url)

    def test_retry(self):
        N = 16
        session = MockSession({
            ('GET', 'uri_1'): (400, 'error!'),
            ('GET', 'uri_2'): [
                (400, 'error!'),
                (400, 'error!'),
                (400, 'error!'),
                (200, 'success'),
            ],
        })
        client = clients.Client(
            retry_policy=policies.BinaryExponentialBackoff(N),
            _session=session,
            _sleep=fake_sleep,
        )

        session._logs.clear()
        with self.assertRaisesRegex(requests.HTTPError, 'http error'):
            client.get('uri_1')
        self.assertEqual(1 + N, len(session._logs))
        for req in session._logs:
            self.assertTrue(isinstance(req, requests.Request))
            self.assertEqual('GET', req.method)
            self.assertEqual('uri_1', req.url)

        session._logs.clear()
        self.assertEqual('success', client.get('uri_2').content)
        self.assertEqual(4, len(session._logs))
        for req in session._logs:
            self.assertTrue(isinstance(req, requests.Request))
            self.assertEqual('GET', req.method)
            self.assertEqual('uri_2', req.url)


class MockSession:

    def __init__(self, req_to_rep):
        self._req_to_rep = req_to_rep<commit_msg>Add unit tests for rate limit behavoir<commit_after>import unittest

import requests

from garage.http2 import clients
from garage.http2 import policies


class TestClient(unittest.TestCase):

    def test_client(self):
        req_to_rep = {
            ('GET', 'uri_1'): (200, 'hello world'),
            ('GET', 'uri_2'): Exception('error_2'),
            ('GET', 'uri_3'): (200, 'GET uri_3'),
            ('HEAD', 'uri_3'): (200, 'HEAD uri_3'),
            ('POST', 'uri_3'): (200, 'POST uri_3'),
            ('PUT', 'uri_3'): (200, 'PUT uri_3'),
            ('GET', 'uri_4'): (404, ''),
        }
        client = clients.Client(
            _session=MockSession(req_to_rep),
            _sleep=fake_sleep,
        )

        rep = client.send(clients.Request('GET', 'uri_1'))
        self.assertEqual('hello world', rep.content)

        with self.assertRaisesRegex(Exception, 'error_2'):
            client.send(clients.Request('GET', 'uri_2'))

        self.assertEqual('GET uri_3', client.get('uri_3').content)
        self.assertEqual('HEAD uri_3', client.head('uri_3').content)
        self.assertEqual('POST uri_3', client.post('uri_3').content)
        self.assertEqual('PUT uri_3', client.put('uri_3').content)

        with self.assertRaisesRegex(requests.HTTPError, 'http error'):
            client.get('uri_4')

    def test_rate_limit(self):
        session = MockSession({
            ('GET', 'uri_1'): [(200, 'tick'), (200, 'tock')],
        })
        client = clients.Client(
            rate_limit=TimedBomb(2),
            _session=session,
            _sleep=fake_sleep,
        )
        self.assertEqual('tick', client.get('uri_1').content)
        self.assertEqual('tock', client.get('uri_1').content)
        with self.assertRaisesRegex(Exception, 'Boom!'):
            client.get('uri_1')

    def test_no_retry(self):
        session = MockSession({('GET', 'uri_1'): (400, 'error!')})
        client = clients.Client(
            retry_policy=policies.NoRetry(),
            _session=session,
            _sleep=fake_sleep,
        )
        with self.assertRaisesRegex(requests.HTTPError, 'http error'):
            client.get('uri_1')
        self.assertEqual(1, len(session._logs))
        for req in session._logs:
            self.assertTrue(isinstance(req, requests.Request))
            self.assertEqual('GET', req.method)
            self.assertEqual('uri_1', req.url)

    def test_retry(self):
        N = 16
        session = MockSession({
            ('GET', 'uri_1'): (400, 'error!'),
            ('GET', 'uri_2'): [
                (400, 'error!'),
                (400, 'error!'),
                (400, 'error!'),
                (200, 'success'),
            ],
        })
        client = clients.Client(
            retry_policy=policies.BinaryExponentialBackoff(N),
            _session=session,
            _sleep=fake_sleep,
        )

        session._logs.clear()
        with self.assertRaisesRegex(requests.HTTPError, 'http error'):
            client.get('uri_1')
        self.assertEqual(1 + N, len(session._logs))
        for req in session._logs:
            self.assertTrue(isinstance(req, requests.Request))
            self.assertEqual('GET', req.method)
            self.assertEqual('uri_1', req.url)

        session._logs.clear()
        self.assertEqual('success', client.get('uri_2').content)
        self.assertEqual(4, len(session._logs))
        for req in session._logs:
            self.assertTrue(isinstance(req, requests.Request))
            self.assertEqual('GET', req.method)
            self.assertEqual('uri_2', req.url)


class TimedBomb:

    def __init__(self, count):
        self.count = count

    def __enter__(self):
        self.count -= 1
        if self.count < 0:
            raise Exception('Boom!')

    def __exit__(self, *_):
        pass


class MockSession:

    def __init__(self, req_to_rep):
        self._req_to_rep = req_to_rep
{'fraction_non_alphanumeric': 0.14805961754780653, 'fraction_numerical': 0.023200224971878514, 'mean_word_length': 2.1925493716337523, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7432391', 'n_tokens_mistral': 2187, 'n_tokens_neox': 2121, 'n_words': 397}