import pytest

from io import BytesIO
from psi_athena.athena import Athena
from test_utils.mock import MockApi
from typing import Tuple


class TestAthena_UnitTests:

    @pytest.mark.unit
    def test_assume_role(self):
        # Set expected values
        expected_role = 'aws:role:arn'
        expected_region = 'region-code'
        expected_session_name = 'test session name'
        expected_access_key = 'accesskey'
        expected_secret_access_key = 'secret'
        expected_session_token = 'token'

        # Set up mocks
        mock_session, _ = self.get_mocks()
        mock_session.on(
            'assume_role',
            RoleArn=expected_role,
            RoleSessionName=expected_session_name
        ).return_val({
            'Credentials': {
                'AccessKeyId': expected_access_key,
                'SecretAccessKey': expected_secret_access_key,
                'SessionToken': expected_session_token
            }
        })
        mock_session.on(
            'client',
            'athena',
            region_name=expected_region,
            aws_access_key_id=expected_access_key,
            aws_secret_access_key=expected_secret_access_key,
            aws_session_token=expected_session_token
        ).return_val(None)

        # Test assume_role()
        a = Athena(aws_session=mock_session)
        a.assume_role(expected_role, expected_region, expected_session_name)
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_run_athena_query_exception(self):
        # Set expected values
        expected_query = 'sql'
        expected_database = 'test_db'
        expected_table = 'test_table'
        expected_catalog = 'TestCatalog'
        expected_workgroup = 'test-wg'
        expected_query_execution_id = 'expected-id'
        expected_bucket = 'bucket'
        expected_key = 'key'
        expected_exception = Exception('test exception')

        # Set up mocks
        mock_session, mock_athena_client = self.get_mocks()
        mock_athena_client.on(
            'start_query_execution',
            QueryString=expected_query,
            QueryExecutionContext={
                'Database': expected_database,
                'Catalog': expected_catalog
            },
            WorkGroup=expected_workgroup
        ).return_val({'QueryExecutionId': expected_query_execution_id})
        mock_athena_client.on(
            'get_query_execution',
            QueryExecutionId=expected_query_execution_id
        ).return_val({
            'QueryExecution': {
                'Status': {'State': 'SUCCEEDED'},
                'ResultConfiguration': {'OutputLocation': f's3://{expected_bucket}/{expected_key}'}
            }
        })
        mock_s3_client = MockApi()
        mock_s3_client.on(
            'get_object',
            Bucket=expected_bucket,
            Key=expected_key
        ).raise_exception(expected_exception)
        mock_session.on('client','s3').return_val(mock_s3_client)

        # Test run_athena_query() exception
        a = Athena(
            data_catalog=expected_catalog,
            workgroup=expected_workgroup,
            database=expected_database,
            aws_session=mock_session
        )
        with pytest.raises(Exception) as e_info:
            a.run_athena_query(expected_table, expected_query, True)
        assert e_info.value == expected_exception
        mock_athena_client.assert_expectations()
        mock_s3_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_start_query(self):
        # Set expected values
        expected_query = 'sql'
        expected_database = 'test_db'
        expected_catalog = 'TestCatalog'
        expected_workgroup = 'test-wg'
        expected_query_execution_id = 'expected-id'

        # Set up mocks
        mock_session, mock_athena_client = self.get_mocks()
        mock_athena_client.on(
            'start_query_execution',
            QueryString=expected_query,
            QueryExecutionContext={
                'Database': expected_database,
                'Catalog': expected_catalog
            },
            WorkGroup=expected_workgroup
        ).return_val({'QueryExecutionId': expected_query_execution_id})

        # Test start_query()
        a = Athena(
            data_catalog=expected_catalog,
            workgroup=expected_workgroup,
            database=expected_database,
            aws_session=mock_session
        )
        query_execution_id = a.start_query(expected_query)
        assert query_execution_id == expected_query_execution_id
        mock_athena_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_start_query_exception(self):
        # Set expected values
        expected_query = 'sql'
        expected_database = 'test_db'
        expected_catalog = 'TestCatalog'
        expected_workgroup = 'test-wg'
        expected_exception = Exception('test')

        # Set up mocks
        mock_session, mock_athena_client = self.get_mocks()
        mock_athena_client.on(
            'start_query_execution',
            QueryString=expected_query,
            QueryExecutionContext={
                'Database': expected_database,
                'Catalog': expected_catalog
            },
            WorkGroup=expected_workgroup
        ).raise_exception(expected_exception)

        # Test start_query() exception
        a = Athena(
            data_catalog=expected_catalog,
            workgroup=expected_workgroup,
            database=expected_database,
            aws_session=mock_session
        )
        with pytest.raises(Exception) as e_info:
            a.start_query(expected_query)
        assert e_info.value == expected_exception
        mock_athena_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_poll_for_query_results(self):
        # Set expected values
        expected_query_execution_id = 'expected-id'
        expected_results = {'QueryExecution': {'Status': {'State': 'SUCCEEDED'}}}

        # Set up mocks
        mock_session, mock_athena_client = self.get_mocks()
        mock_athena_client.on(
            'get_query_execution',
            QueryExecutionId=expected_query_execution_id
        ).return_val(expected_results)

        # Test poll_for_query_results()
        a = Athena(aws_session=mock_session)
        results = a.poll_for_query_results(expected_query_execution_id)
        assert results == expected_results
        mock_athena_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_poll_for_query_results_exception(self):
        # Set expected values
        expected_query_execution_id = 'expected-id'
        expected_exception = Exception('test')

        # Set up mocks
        mock_session, mock_athena_client = self.get_mocks()
        mock_athena_client.on(
            'get_query_execution',
            QueryExecutionId=expected_query_execution_id
        ).raise_exception(expected_exception)

        # Test poll_for_query_results() exception
        a = Athena(aws_session=mock_session)
        with pytest.raises(Exception) as e_info:
            a.poll_for_query_results(expected_query_execution_id)
        assert e_info.value == expected_exception
        mock_athena_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_download_query_results_as_dict(self):
        # Set expected values
        expected_bucket = 'bucket'
        expected_key = 'key'
        expected_body = 'test,csv,body\nalpha.val,,4.0\nbeta.val,-10,'.encode('utf-8')

        # Set up mocks
        mock_s3_client = MockApi()
        mock_s3_client.on(
            'get_object',
            Bucket=expected_bucket,
            Key=expected_key
        ).return_val({'Body': BytesIO(expected_body)})
        mock_session, _ = self.get_mocks()
        mock_session.on('client','s3').return_val(mock_s3_client)

        # Test download_query_results_as_dict()
        a = Athena(aws_session=mock_session)
        results = a.download_query_results_as_dict({
            'QueryExecution': {
                'ResultConfiguration': {
                    'OutputLocation': f's3://{expected_bucket}/{expected_key}'
                }
            }
        })
        assert results == [{
            'test': 'alpha.val',
            'csv': None,
            'body': 4.0
        }, {
            'test': 'beta.val',
            'csv': -10,
            'body': None
        }]
        mock_s3_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_download_query_results_as_dict_exception(self):
        # Set expected values
        expected_bucket = 'bucket'
        expected_key = 'key'
        expected_body = 'test,csv,body\nalpha.val,-10,4.0,extra,columns'.encode('utf-8')

        # Set up mocks
        mock_s3_client = MockApi()
        mock_s3_client.on(
            'get_object',
            Bucket=expected_bucket,
            Key=expected_key
        ).return_val({'Body': BytesIO(expected_body)})
        mock_session, _ = self.get_mocks()
        mock_session.on('client','s3').return_val(mock_s3_client)

        # Test download_query_results_as_dict()
        a = Athena(aws_session=mock_session)
        with pytest.raises(AssertionError) as e_info:
            a.download_query_results_as_dict({
                'QueryExecution': {
                    'ResultConfiguration': {
                        'OutputLocation': f's3://{expected_bucket}/{expected_key}'
                    }
                }
            })
        assert str(e_info.value).startswith('Invalid CSV column')
        mock_s3_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_download_query_results(self):
        # Set expected values
        expected_bucket = 'bucket'
        expected_key = 'key'
        expected_body = 'test,csv,body'.encode('utf-8')

        # Set up mocks
        mock_s3_client = MockApi()
        mock_s3_client.on(
            'get_object',
            Bucket=expected_bucket,
            Key=expected_key
        ).return_val({'Body': BytesIO(expected_body)})
        mock_session, _ = self.get_mocks()
        mock_session.on('client','s3').return_val(mock_s3_client)

        # Test download_query_results()
        a = Athena(aws_session=mock_session)
        results = a.download_query_results({
            'QueryExecution': {
                'ResultConfiguration': {
                    'OutputLocation': f's3://{expected_bucket}/{expected_key}'
                }
            }
        })
        assert results == expected_body
        mock_s3_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_download_query_results_exception(self):
        # Set expected values
        expected_bucket = 'bucket'
        expected_key = 'key'
        expected_exception = Exception('test exception')

        # Set up mocks
        mock_s3_client = MockApi()
        mock_s3_client.on(
            'get_object',
            Bucket=expected_bucket,
            Key=expected_key
        ).raise_exception(expected_exception)
        mock_session, _ = self.get_mocks()
        mock_session.on('client','s3').return_val(mock_s3_client)

        # Test download_query_results() exception
        a = Athena(aws_session=mock_session)
        with pytest.raises(Exception) as e_info:
            a.download_query_results({
                'QueryExecution': {
                    'ResultConfiguration': {
                        'OutputLocation': f's3://{expected_bucket}/{expected_key}'
                    }
                }
            })
        assert e_info.value == expected_exception
        mock_s3_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_reload_partitions_exception(self):
        # Set expected values
        expected_table = 'test_table'
        expected_exception = Exception('test')

        # Set up mocks
        mock_session, mock_athena_client = self.get_mocks()
        mock_athena_client.on(
            'start_query_execution',
            QueryString=f'MSCK REPAIR TABLE `{expected_table}`',
            QueryExecutionContext={
                'Database': 'default',
                'Catalog': 'AwsDataCatalog'
            },
            WorkGroup='primary'
        ).raise_exception(expected_exception)

        # Test reload_partitions() exception
        a = Athena(aws_session=mock_session)
        with pytest.raises(Exception) as e_info:
            a.reload_partitions(expected_table)
        assert e_info.value == expected_exception
        mock_athena_client.assert_expectations()
        mock_session.assert_expectations()

    def get_mocks(self) -> Tuple[MockApi, MockApi]:
        mock_athena_client = MockApi()
        mock_session = MockApi()
        mock_session.on('client','athena').return_val(mock_athena_client)
        return mock_session, mock_athena_client
