import pytest
import pytz

from botocore.exceptions import ClientError
from datetime import datetime
from psi_kinesis.kinesis_reader import KinesisReader, PROVISIONED_THROUGHPUT_EXCEPTION
from psi_utils.unittest_mock import MockApi
from typing import Tuple


class TestKinesisReader_UnitTests:

    @pytest.mark.unit
    def test_assume_role(self):
        # Set expected values
        expected_role = "RoleName"
        expected_account = "01234567890"
        expected_role_arn = "arn:aws:iam::01234567890:role/RoleName"
        expected_role_session_name = "RoleSession"
        expected_err = Exception("test-err")

        # Set up mocks
        mock_sts_client = MockApi()
        mock_session, _ = self.get_mocks()
        mock_session.on("client", "sts").return_val(mock_sts_client)
        mock_sts_client.on(
            "assume_role", RoleArn=expected_role_arn, RoleSessionName=expected_role_session_name
        ).raise_exception(expected_err)

        # Test assume_role() Exception
        kr = KinesisReader(aws_session=mock_session)
        with pytest.raises(Exception) as e_info:
            kr.assume_role(
                expected_role, expected_account, role_session_name=expected_role_session_name
            )
        assert e_info.value == expected_err
        mock_sts_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_does_stream_exist(self):
        # Set expected values
        expected_stream_name = "kinesis-stream-name"
        expected_resp = {"StreamDescription": {"StreamName": expected_stream_name}}

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on("describe_stream", StreamName=expected_stream_name).return_val(
            expected_resp
        )

        # Test does_stream_exist()
        kr = KinesisReader(aws_session=mock_session)
        assert kr.does_stream_exist(expected_stream_name)
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_does_stream_exist_resource_not_found(self):
        # Set expected values
        expected_stream_name = "kinesis-stream-name"
        expected_err = ClientError({"Error": {"Code": "ResourceNotFoundException"}}, "test")

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on("describe_stream", StreamName=expected_stream_name).raise_exception(
            expected_err
        )

        # Test does_stream_exist() ResourceNotFoundException
        kr = KinesisReader(aws_session=mock_session)
        assert not kr.does_stream_exist(expected_stream_name)
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_does_stream_exist_client_error(self):
        # Set expected values
        expected_stream_name = "kinesis-stream-name"
        expected_err = ClientError({"Error": {"Code": "ClientError"}}, "test")

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on("describe_stream", StreamName=expected_stream_name).raise_exception(
            expected_err
        )

        # Test does_stream_exist() ClientError
        kr = KinesisReader(aws_session=mock_session)
        with pytest.raises(ClientError) as e_info:
            kr.does_stream_exist(expected_stream_name)
        assert e_info.value == expected_err
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_does_stream_exist_error(self):
        # Set expected values
        expected_stream_name = "kinesis-stream-name"
        expected_err = Exception("test-exception")

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on("describe_stream", StreamName=expected_stream_name).raise_exception(
            expected_err
        )

        # Test does_stream_exist() Exception
        kr = KinesisReader(aws_session=mock_session)
        with pytest.raises(Exception) as e_info:
            kr.does_stream_exist(expected_stream_name)
        assert e_info.value == expected_err
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_list_shards(self):
        # Set expected values
        expected_stream_name = "kinesis-stream-name"
        expected_shard_filter = {
            "Type": "FROM_TIMESTAMP",
            "Timestamp": datetime(2022, 4, 1, tzinfo=pytz.UTC),
        }
        expected_shards = [{"ShardId": "shard-id-000001"}]

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            "list_shards", StreamName=expected_stream_name, ShardFilter=expected_shard_filter
        ).return_val({"Shards": expected_shards})

        # Test list_shards()
        kr = KinesisReader(aws_session=mock_session)
        shards = kr.list_shards(expected_stream_name, shard_filter=expected_shard_filter)
        assert shards == expected_shards
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_list_shards_error(self):
        # Set expected values
        expected_stream_name = "kinesis-stream-name"
        expected_err = Exception("test-exception")

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on("list_shards", StreamName=expected_stream_name).raise_exception(
            expected_err
        )

        # Test list_shards() Exception
        kr = KinesisReader(aws_session=mock_session)
        with pytest.raises(Exception) as e_info:
            kr.list_shards(expected_stream_name)
        assert e_info.value == expected_err
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_get_records(self):
        # Set expected values
        expected_stream_name = "kinesis-stream-name"
        expected_shard_id = "shard-id-000001"
        expected_shard_iterator = "shard-iterator-id"
        expected_records = [{"a": 1}, {"b": 2}, {"c": 3}]

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            "get_shard_iterator",
            StreamName=expected_stream_name,
            ShardId=expected_shard_id,
            ShardIteratorType="TRIM_HORIZON",
        ).return_val({"ShardIterator": expected_shard_iterator})
        mock_kinesis_client.on(
            "get_records", ShardIterator=expected_shard_iterator, Limit=10000
        ).return_val({"Records": expected_records, "MillisBehindLatest": 0})

        # Test get_records() with a sequence number but no shard iterator type
        kr = KinesisReader(aws_session=mock_session)
        with pytest.raises(AssertionError) as e_info:
            kr.get_records(
                expected_stream_name,
                shard_id=expected_shard_id,
                starting_sequence_number="1234567890123456789",
            )
        assert (
            str(e_info.value)
            == "Cannot specify a starting sequence number or timestamp if shard iteratory type is one of [TRIM_HORIZON,LATEST]"
        )

        # Test get_records() with a timestamp but no shard iterator type
        kr = KinesisReader(aws_session=mock_session)
        with pytest.raises(AssertionError) as e_info:
            kr.get_records(
                expected_stream_name,
                shard_id=expected_shard_id,
                timestamp=datetime(2022, 4, 1, tzinfo=pytz.UTC),
            )
        assert (
            str(e_info.value)
            == "Cannot specify a starting sequence number or timestamp if shard iteratory type is one of [TRIM_HORIZON,LATEST]"
        )

        # Test get_records()
        kr = KinesisReader(aws_session=mock_session)
        records = kr.get_records(expected_stream_name, shard_id=expected_shard_id)
        assert records == {expected_shard_id: expected_records}
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_get_records_at_sequence_number(self):
        # Set expected values
        expected_stream_name = "kinesis-stream-name"
        expected_shard_id = "shard-id-000001"
        expected_shard_iterator = "shard-iterator-id"
        expected_shard_iterator_type = "AT_SEQUENCE_NUMBER"
        expected_sequence_number = "1234567890123456789"
        expected_records = [{"a": 1}, {"b": 2}, {"c": 3}]

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            "get_shard_iterator",
            StreamName=expected_stream_name,
            ShardId=expected_shard_id,
            ShardIteratorType=expected_shard_iterator_type,
            StartingSequenceNumber=expected_sequence_number,
        ).return_val({"ShardIterator": expected_shard_iterator})
        mock_kinesis_client.on(
            "get_records", ShardIterator=expected_shard_iterator, Limit=10000
        ).return_val({"Records": expected_records, "MillisBehindLatest": 0})

        # Test get_records() AT_SEQUENCE_NUMBER
        kr = KinesisReader(aws_session=mock_session)
        records = kr.get_records(
            expected_stream_name,
            shard_id=expected_shard_id,
            shard_iterator_type=expected_shard_iterator_type,
            starting_sequence_number=expected_sequence_number,
        )
        assert records == {expected_shard_id: expected_records}
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_get_records_at_timestamp(self):
        # Set expected values
        expected_stream_name = "kinesis-stream-name"
        expected_shard_id = "shard-id-000001"
        expected_shard_iterator = "shard-iterator-id"
        expected_shard_iterator_type = "AT_TIMESTAMP"
        expected_timestamp = datetime(2022, 4, 1, tzinfo=pytz.UTC)
        expected_records = [{"a": 1}, {"b": 2}, {"c": 3}]

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            "get_shard_iterator",
            StreamName=expected_stream_name,
            ShardId=expected_shard_id,
            ShardIteratorType=expected_shard_iterator_type,
            Timestamp=expected_timestamp,
        ).return_val({"ShardIterator": expected_shard_iterator})
        mock_kinesis_client.on(
            "get_records", ShardIterator=expected_shard_iterator, Limit=10000
        ).return_val({"Records": expected_records, "MillisBehindLatest": 0})

        # Test get_records() AT_TIMESTAMP
        kr = KinesisReader(aws_session=mock_session)
        records = kr.get_records(
            expected_stream_name,
            shard_id=expected_shard_id,
            shard_iterator_type=expected_shard_iterator_type,
            timestamp=expected_timestamp,
        )
        assert records == {expected_shard_id: expected_records}
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_get_records_for_entire_stream(self):
        # Set expected values
        expected_stream_name = "kinesis-stream-name"
        expected_shard_iterator = "shard-iterator-id"
        expected_records = [{"a": 1}, {"b": 2}, {"c": 3}]
        expected_shard_id = "shard-id-000001"

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on("list_shards", StreamName=expected_stream_name).return_val(
            {"Shards": [{"ShardId": expected_shard_id}]}
        )
        mock_kinesis_client.on(
            "get_shard_iterator",
            StreamName=expected_stream_name,
            ShardId=expected_shard_id,
            ShardIteratorType="TRIM_HORIZON",
        ).return_val({"ShardIterator": expected_shard_iterator})
        mock_kinesis_client.on(
            "get_records", ShardIterator=expected_shard_iterator, Limit=10000
        ).return_val({"Records": expected_records, "MillisBehindLatest": 0})

        # Test get_records() for an entire Kinesis stream
        kr = KinesisReader(aws_session=mock_session)
        records = kr.get_records(expected_stream_name)
        assert records == {expected_shard_id: expected_records}
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_get_records_with_shard_iterator(self):
        # Set expected values
        fst_shard_iterator = "first-shard-iterator-id"
        snd_shard_iterator = "second-shard-iterator-id"
        fst_records = [{"a": 1}, {"b": 2}]
        snd_records = [{"c": 3}, {"d": 4}]
        fst_resp = {
            "Records": fst_records,
            "NextShardIterator": snd_shard_iterator,
            "MillisBehindLatest": 123,
        }
        snd_resp = {
            "Records": snd_records,
            "NextShardIterator": "third-shard-iterator-id",
            "MillisBehindLatest": 0,
        }
        expected_records = fst_records + snd_records

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            "get_records", ShardIterator=fst_shard_iterator, Limit=10000
        ).return_val(fst_resp)
        mock_kinesis_client.on(
            "get_records", ShardIterator=snd_shard_iterator, Limit=10000
        ).return_val(snd_resp)

        # Test get_records_with_shard_iterator()
        kr = KinesisReader(aws_session=mock_session)
        records = kr.get_records_with_shard_iterator(fst_shard_iterator)
        assert records == expected_records
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__get_shard_iterator(self):
        # Set expected values
        expected_stream_name = "kinesis-stream-name"
        expected_shard_id = "shard-id-000001"
        expected_shard_iterator_type = "AT_SEQUENCE_NUMBER"
        expected_sequence_number = "1234567890123456789"
        expected_shard_iterator = "shard-iterator-id"
        expected_resp = {"ShardIterator": expected_shard_iterator}

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            "get_shard_iterator",
            StreamName=expected_stream_name,
            ShardId=expected_shard_id,
            ShardIteratorType=expected_shard_iterator_type,
            StartingSequenceNumber=expected_sequence_number,
        ).return_val(expected_resp)

        # Test __get_shard_iterator()
        kr = KinesisReader(aws_session=mock_session)
        shard_iterator = kr._KinesisReader__get_shard_iterator(
            expected_stream_name,
            expected_shard_id,
            expected_shard_iterator_type,
            starting_sequence_number=expected_sequence_number,
        )
        assert shard_iterator == expected_shard_iterator
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__get_shard_iterator_throughput_exception(self):
        # Set expected values
        expected_stream_name = "kinesis-stream-name"
        expected_shard_id = "shard-id-000001"
        expected_shard_iterator_type = "AT_TIMESTAMP"
        expected_timestamp = datetime(2022, 4, 1, tzinfo=pytz.UTC)
        expected_err = ClientError({"Error": {"Code": PROVISIONED_THROUGHPUT_EXCEPTION}}, "test")

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            "get_shard_iterator",
            StreamName=expected_stream_name,
            ShardId=expected_shard_id,
            ShardIteratorType=expected_shard_iterator_type,
            Timestamp=expected_timestamp,
        ).raise_exception(expected_err)

        # Test __get_shard_iterator() ProvisionedThroughputExceededException
        kr = KinesisReader(aws_session=mock_session, max_retry_count=1)

        with pytest.raises(AssertionError) as e_info:
            kr._KinesisReader__get_shard_iterator(
                expected_stream_name,
                expected_shard_id,
                expected_shard_iterator_type,
                timestamp=expected_timestamp,
            )
        assert str(e_info.value) == "Retries have been exhausted"
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__get_shard_iterator_client_error(self):
        # Set expected values
        expected_stream_name = "kinesis-stream-name"
        expected_shard_id = "shard-id-000001"
        expected_shard_iterator_type = "AT_TIMESTAMP"
        expected_timestamp = datetime(2022, 4, 1, tzinfo=pytz.UTC)
        expected_err = ClientError({"Error": {"Code": "ClientError"}}, "test")

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            "get_shard_iterator",
            StreamName=expected_stream_name,
            ShardId=expected_shard_id,
            ShardIteratorType=expected_shard_iterator_type,
            Timestamp=expected_timestamp,
        ).raise_exception(expected_err)

        # Test __get_shard_iterator() ClientError
        kr = KinesisReader(aws_session=mock_session, max_retry_count=1)

        with pytest.raises(ClientError) as e_info:
            kr._KinesisReader__get_shard_iterator(
                expected_stream_name,
                expected_shard_id,
                expected_shard_iterator_type,
                timestamp=expected_timestamp,
            )
        assert e_info.value == expected_err
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__get_shard_iterator_error(self):
        # Set expected values
        expected_stream_name = "kinesis-stream-name"
        expected_shard_id = "shard-id-000001"
        expected_shard_iterator_type = "AT_TIMESTAMP"
        expected_timestamp = datetime(2022, 4, 1, tzinfo=pytz.UTC)
        expected_err = Exception("test-exception")

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            "get_shard_iterator",
            StreamName=expected_stream_name,
            ShardId=expected_shard_id,
            ShardIteratorType=expected_shard_iterator_type,
            Timestamp=expected_timestamp,
        ).raise_exception(expected_err)

        # Test __get_shard_iterator() Exception
        kr = KinesisReader(aws_session=mock_session, max_retry_count=1)

        with pytest.raises(Exception) as e_info:
            kr._KinesisReader__get_shard_iterator(
                expected_stream_name,
                expected_shard_id,
                expected_shard_iterator_type,
                timestamp=expected_timestamp,
            )
        assert e_info.value == expected_err
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__get_records_with_shard_iterator(self):
        # Set expected values
        expected_shard_iterator = "shard-iterator-id"
        expected_resp = {"Records": [], "NextShardIterator": None, "MillisBehindLatest": 123}

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            "get_records", ShardIterator=expected_shard_iterator, Limit=10000
        ).return_val(expected_resp)

        # Test __get_records_with_shard_iterator()
        kr = KinesisReader(aws_session=mock_session)
        resp = kr._KinesisReader__get_records_with_shard_iterator(expected_shard_iterator)
        assert resp == expected_resp
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__get_records_with_shard_iterator_throughput_exception(self):
        # Set expected values
        expected_shard_iterator = "shard-iterator-id"
        expected_err = ClientError({"Error": {"Code": PROVISIONED_THROUGHPUT_EXCEPTION}}, "test")

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            "get_records", ShardIterator=expected_shard_iterator, Limit=10000
        ).raise_exception(expected_err)

        # Test __get_records_with_shard_iterator() ProvisionedThroughputExceededException
        kr = KinesisReader(aws_session=mock_session, max_retry_count=1)

        with pytest.raises(AssertionError) as e_info:
            kr._KinesisReader__get_records_with_shard_iterator(expected_shard_iterator)
        assert str(e_info.value) == "Retries have been exhausted"
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__get_records_with_shard_iterator_client_error(self):
        # Set expected values
        expected_shard_iterator = "shard-iterator-id"
        expected_err = ClientError({"Error": {"Code": "ClientError"}}, "test")

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            "get_records", ShardIterator=expected_shard_iterator, Limit=10000
        ).raise_exception(expected_err)

        # Test __get_records_with_shard_iterator() ClientError
        kr = KinesisReader(aws_session=mock_session, max_retry_count=1)

        with pytest.raises(ClientError) as e_info:
            kr._KinesisReader__get_records_with_shard_iterator(expected_shard_iterator)
        assert e_info.value == expected_err
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__get_records_with_shard_iterator_error(self):
        # Set expected values
        expected_shard_iterator = "shard-iterator-id"
        expected_err = Exception("test exception")

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            "get_records", ShardIterator=expected_shard_iterator, Limit=10000
        ).raise_exception(expected_err)

        # Test __get_records_with_shard_iterator() Exception
        kr = KinesisReader(aws_session=mock_session, max_retry_count=1)

        with pytest.raises(Exception) as e_info:
            kr._KinesisReader__get_records_with_shard_iterator(expected_shard_iterator)
        assert e_info.value == expected_err
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_next_wait_seconds(self):
        retry_backoff_rate = 2.5
        mock_session, _ = self.get_mocks()
        kr = KinesisReader(aws_session=mock_session, retry_backoff_rate=retry_backoff_rate)
        assert kr._next_wait_seconds(None) == 1.0
        assert kr._next_wait_seconds(1.0) == retry_backoff_rate
        assert kr._next_wait_seconds(retry_backoff_rate) == 6.25
        assert kr._next_wait_seconds(10) == 25

    @pytest.mark.unit
    def test_next_retry_count(self):
        mock_session, _ = self.get_mocks()
        kr = KinesisReader(aws_session=mock_session)
        assert kr._next_retry_count(None) == 1
        assert kr._next_retry_count(10) == 11

    def get_mocks(self) -> Tuple[MockApi, MockApi]:
        mock_kinesis_client = MockApi()
        mock_session = MockApi()
        mock_session.on("client", "kinesis").return_val(mock_kinesis_client)
        return mock_session, mock_kinesis_client
