"""Athena test module."""
import pytest
from moto import mock_aws

from src.utils import athena


@pytest.fixture()
def athena_instance(aws_credentials):
    """Return an AthenaUtil instance."""
    with mock_aws():
        athena_util_instance = athena.AthenaUtil()
        yield athena_util_instance


@pytest.fixture()
def mock_output_location():
    """Return a test output location."""
    return 's3://dev-bucket/athena/'


@pytest.fixture()
def mock_list_named_queries(athena_instance):
    """Return a list of named queries."""
    list_named_queries = athena_instance.athena_client.list_named_queries = lambda WorkGroup: {  # noqa: E501
        'NamedQueryIds': [
            'query1',
            'query2',
        ]
    }
    return list_named_queries


@pytest.fixture()
def mock_get_named_query(athena_instance, mock_list_named_queries):
    """Return a named query."""
    get_named_query = athena_instance.athena_client.get_named_query = lambda NamedQueryId: {  # noqa: E501
        'NamedQuery': {
            'Name': 'shared_1234567890_break_glass',
            'QueryString': f'SELECT * FROM table_{NamedQueryId}'
        }
    }
    return get_named_query


@pytest.fixture()
def mock_start_query_execution(athena_instance):
    """Return a query execution."""
    start_query_execution = athena_instance.athena_client.start_query_execution = lambda QueryString, QueryExecutionContext, WorkGroup, ExecutionParameters: {  # noqa: E501
        'QueryExecutionId': 'execution1',
        'QueryString': 'SELECT * FROM table_query1',
        'QueryExecutionContext': {
            'Database': 'shared_1234567890_break_glass',
            'Catalog': '1234567890'
        },
        'WorkGroup': 'shareed_break_glass',
        'ExecutionParameters': [
            '2023/12/11',
            '2023/12/11',
            'dev-break-glass-access-role',
            'test_user'
        ]
    }
    return start_query_execution


@pytest.fixture()
def mock_get_query_execution(athena_instance, mock_output_location):
    """Return a query execution."""
    get_query_execution = athena_instance.athena_client.get_query_execution = lambda QueryExecutionId: {  # noqa: E501
        'QueryExecution': {
            'Status': {
                'State': 'SUCCEEDED'
            },
            'ResultConfiguration': {
                'OutputLocation': f'{mock_output_location}/output.csv'
            }
        }
    }
    return get_query_execution


@pytest.fixture()
def mock_get_failed_query_execution(athena_instance):
    """Return a failed query execution."""
    get_query_execution = athena_instance.athena_client.get_query_execution = lambda QueryExecutionId: {  # noqa: E501
        'QueryExecution': {
            'Status': {
                'State': 'FAILED',
                'StateChangeReason': 'Query failed',
                'AthenaError': {
                    'ErrorCode': '123',
                    'ErrorMessage': 'Query failed'
                }
            }
        }
    }
    return get_query_execution


def test_get_named_query_statements(athena_instance, mock_get_named_query):
    """Test get_named_query_statements function."""
    result = athena_instance.get_named_query_statements('1234567890')
    assert result is not None
    assert result == 'SELECT * FROM table_query1'


def test_get_query_execution_details(athena_instance, mock_output_location, mock_get_query_execution):  # noqa: E501
    """Test get_query_execution_details function."""
    result = athena_instance.get_query_execution_details('query1')
    assert result is not None
    assert result == f'{mock_output_location}/output.csv'


def test_get_query_execution_details_failed(athena_instance, mock_get_failed_query_execution):  # noqa: E501
    """Test get_query_execution_details function."""
    with pytest.raises(SystemExit):
        athena_instance.get_query_execution_details('query1')


def test_start_audit_query_execution(athena_instance, mock_start_query_execution):  # noqa: E501
    """Test start_audit_query_execution function."""
    result = athena_instance.start_audit_query_execution(
        'test_user', '2023/12/12', 'SELECT * FROM table_query1')
    assert result is not None
    assert result == 'execution1'