"""Test athena."""
from unittest.mock import patch

import pytest

from feed_ingestion.util.aws import athena

_SUCCEEDED = {'QueryExecution': {'Status': {'State': 'SUCCEEDED'}}}
_PENDING = {'QueryExecution': {'Status': {'State': 'pending'}}}
_S3_BUCKET = 'dev-cucumbers'
_S3_PATH = 'dev-efedorov/athena/result/swf_unit_test/04/'
_QUERY = 'select * from "efedorov"."content_owner_asset_a2" limit 10'


@patch.object(athena, 'boto3')
@patch.object(athena, 'time')
def test_query_athena(time_mock, boto3_mock):
    """Test run_query executes CREATE TABLE and DROP TABLE queries."""
    time_mock.time.return_value = 123
    athena_client_mock = boto3_mock.Session.return_value.client.return_value
    # pretend 2 iterations of pending and then SUCCEEDED
    # second SUCCEEDED required for DROP table statement
    athena_client_mock.get_query_execution.side_effect = [
        _PENDING, _PENDING, _SUCCEEDED, _SUCCEEDED]

    athena.run_query(
        athena_query=_QUERY,
        athena_temp_database='efedorov',
        athena_workgroup='primary',
        destination_s3_bucket=_S3_BUCKET,
        destination_s3_path=_S3_PATH,
    )

    assert time_mock.sleep.call_count == 2
    queries_mock = athena_client_mock.start_query_execution
    assert queries_mock.call_count == 2
    assert 'CREATE TABLE' in queries_mock.call_args_list[0][1]['QueryString']
    assert 'DROP TABLE' in queries_mock.call_args_list[1][1]['QueryString']


@patch.object(athena, 'boto3')
@patch.object(athena, 'time')
def test_query_athena_timeout(time_mock, boto3_mock):
    """Test run_query raises TimeoutError and still runs DROP TABLE."""
    time_initial = 123
    timeout = 30 * 60
    time_timed_out = time_initial + timeout + 1
    time_mock.time.side_effect = (
        [time_initial, time_initial + 333, time_timed_out]
        + [1, 1]  # for DROP table polling
    )

    athena_client_mock = boto3_mock.Session.return_value.client.return_value
    athena_client_mock.get_query_execution.side_effect = [
        _PENDING, _PENDING, _SUCCEEDED]

    with pytest.raises(TimeoutError):
        athena.run_query(
            athena_query=_QUERY,
            athena_temp_database='efedorov',
            athena_workgroup='primary',
            destination_s3_bucket=_S3_BUCKET,
            destination_s3_path=_S3_PATH,
        )

    queries_mock = athena_client_mock.start_query_execution
    assert queries_mock.call_count == 2
    assert 'CREATE TABLE' in queries_mock.call_args_list[0][1]['QueryString']
    assert 'DROP TABLE' in queries_mock.call_args_list[1][1]['QueryString']


@patch.object(athena, 'boto3')
@pytest.mark.parametrize('state', ['FAILED', 'CANCELLED'])
def test_query_athena_negative(boto3_mock, state):
    """Test run_query raises StatusError and still runs DROP TABLE."""
    athena_client_mock = boto3_mock.Session.return_value.client.return_value
    athena_client_mock.get_query_execution.side_effect = [
        {'QueryExecution': {'Status': {'State': state}}},
        _SUCCEEDED,
    ]

    with pytest.raises(athena.StatusError):
        athena.run_query(
            athena_query=_QUERY,
            athena_temp_database='efedorov',
            athena_workgroup='primary',
            destination_s3_bucket=_S3_BUCKET,
            destination_s3_path=_S3_PATH,
        )

    queries_mock = athena_client_mock.start_query_execution
    assert queries_mock.call_count == 2
    assert 'CREATE TABLE' in queries_mock.call_args_list[0][1]['QueryString']
    # Still should clear the table
    assert 'DROP TABLE' in queries_mock.call_args_list[1][1]['QueryString']


@patch.object(athena, 'boto3')
def test_query_athena_use_unload_query(boto3_mock):
    """Test run_query with use_unload_query=True runs one UNLOAD query."""
    athena_client_mock = boto3_mock.Session.return_value.client.return_value
    athena_client_mock.get_query_execution.return_value = _SUCCEEDED

    athena.run_query(
        athena_query=_QUERY,
        athena_temp_database='',
        athena_workgroup='primary',
        destination_s3_bucket=_S3_BUCKET,
        destination_s3_path=_S3_PATH,
        use_unload_query=True,
    )

    queries_mock = athena_client_mock.start_query_execution
    assert queries_mock.call_count == 1
    assert 'UNLOAD' in queries_mock.call_args_list[0][1]['QueryString']


@patch.object(athena, 'assumed_session')
def test_query_athena_assume_role(assumed_session_mock):
    """Test run_query with assume_role uses assumed_session for the client."""
    role_arn = 'arn:aws:iam::123456789:role/test-role'
    athena_client_mock = (
        assumed_session_mock.return_value.client.return_value)
    athena_client_mock.get_query_execution.return_value = _SUCCEEDED

    athena.run_query(
        athena_query=_QUERY,
        athena_temp_database='',
        athena_workgroup='primary',
        destination_s3_bucket=_S3_BUCKET,
        destination_s3_path=_S3_PATH,
        use_unload_query=True,
        assume_role=role_arn,
    )

    assumed_session_mock.assert_called_once_with(role_arn)
    athena_client_mock.start_query_execution.assert_called_once()