from unittest.mock import Mock

import pytest
from stream.exceptions import RateLimitReached

from getstream_connector.utils.retry import on_ratelimit_reached


# Mock StreamClient class with ratelimit_info attribute
class MockStreamClient:
    class RateLimitInfo:
        def __init__(self, reset_in_seconds):
            self.reset_in_seconds = reset_in_seconds

    def __init__(self, reset_in_seconds):
        self.ratelimit_info = self.RateLimitInfo(reset_in_seconds)


def test_on_ratelimit_reached_success():
    client = MockStreamClient(reset_in_seconds=0.1)

    @on_ratelimit_reached
    def dummy_method(self, *args, **kwargs):
        return "success"

    result = dummy_method(client)
    assert result == "success"


def test_on_ratelimit_reached_with_retries():
    client = MockStreamClient(reset_in_seconds=0.1)
    mock = Mock(
        side_effect=(
            RateLimitReached("message_1"),
            RateLimitReached("message_2"),
            "success",
        )
    )

    @on_ratelimit_reached
    def dummy_method(self, *args, **kwargs):
        return mock()

    result = dummy_method(client)
    assert result == "success"
    assert mock.call_count == 3


def test_on_ratelimit_reached_exceeds_max_retries():
    client = MockStreamClient(reset_in_seconds=0.1)

    mock = Mock(
        side_effect=(
            RateLimitReached("message_1"),
            RateLimitReached("message_2"),
            RateLimitReached("message_3"),
        )
    )

    @on_ratelimit_reached
    def dummy_method(self, *args, **kwargs):
        return mock()

    with pytest.raises(RateLimitReached):
        dummy_method(client)
    assert mock.call_count == 3