# Copyright (c) 2020 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

import pytest

from czc.exceptions import CloudWatchQueryException
from czt.unittest import Context, MockEx


@pytest.fixture(scope='function')
def context(mocker):
    import czc.cloudwatch as cloudwatch

    with Context(mocker, cloudwatch, ['aws_service', 'sleep'], extra={'cloudwatch': cloudwatch}) as context:
        yield context


TEST_LOG_GROUP_NAME = 'log-group-name'
TEST_QUERY_STRING = 'find stuff | filter it | sort it'
TEST_START_TIME = 1648666976
TEST_END_TIME = 1648866976
TEST_QUERY_ID = 'test-query-id'

TEST_RUNNING_QUERY = {'status': 'Running'}
TEST_UNKNOWN_QUERY = {'status': 'Unknown'}
TEST_FAILED_QUERY = {'status': 'Failed'}
TEST_CANCELLED_QUERY = {'status': 'Cancelled'}
TEST_TIMEOUT_QUERY = {'status': 'Timeout'}

TEST_COMPLETED_QUERY = {
    'status': 'Complete',
    'results': [
        [
            {'field': '@timestamp', 'value': 1648666976},
            {'field': '@message', 'value': 'Here is a message'},
        ],
        [
            {'field': '@timestamp', 'value': 1648666977},
            {'field': '@message', 'value': 'Here is another message'},
            {'field': 'something_else', 'value': 'Random field value'},
        ],
    ],
}


@pytest.mark.unit
@pytest.mark.parametrize('iterations', [0, 2, 5])
def test_query_success(context, iterations):
    MockEx(context.mock_aws_service.logs_client.start_query).if_called_with(
        logGroupName=TEST_LOG_GROUP_NAME,
        queryString=TEST_QUERY_STRING,
        startTime=TEST_START_TIME,
        endTime=TEST_END_TIME,
        limit=100,
    ).return_value({'queryId': TEST_QUERY_ID})

    def get_query_results_call(*args, **kwargs):
        return (
            TEST_COMPLETED_QUERY
            if context.mock_aws_service.logs_client.get_query_results.call_count == (iterations + 1)
            else TEST_RUNNING_QUERY
        )

    context.mock_aws_service.logs_client.get_query_results.side_effect = get_query_results_call

    results = context.cloudwatch.query(TEST_LOG_GROUP_NAME, TEST_QUERY_STRING, TEST_START_TIME, TEST_END_TIME)

    assert results == [
        {'@timestamp': 1648666976, '@message': 'Here is a message'},
        {'@timestamp': 1648666977, '@message': 'Here is another message', 'something_else': 'Random field value'},
    ]

    assert MockEx(context.mock_aws_service.logs_client.get_query_results).was_called_with(queryId=TEST_QUERY_ID)


@pytest.mark.unit
@pytest.mark.parametrize('iterations', [0, 2, 5])
@pytest.mark.parametrize('result', [TEST_FAILED_QUERY, TEST_CANCELLED_QUERY, TEST_TIMEOUT_QUERY])
def test_query_failure(context, iterations, result):
    MockEx(context.mock_aws_service.logs_client.start_query).if_called_with(
        logGroupName=TEST_LOG_GROUP_NAME,
        queryString=TEST_QUERY_STRING,
        startTime=TEST_START_TIME,
        endTime=TEST_END_TIME,
        limit=100,
    ).return_value({'queryId': TEST_QUERY_ID})

    def get_query_results_call(*args, **kwargs):
        return (
            result
            if context.mock_aws_service.logs_client.get_query_results.call_count == (iterations + 1)
            else TEST_RUNNING_QUERY
        )

    context.mock_aws_service.logs_client.get_query_results.side_effect = get_query_results_call

    with pytest.raises(CloudWatchQueryException) as exc:
        context.cloudwatch.query(TEST_LOG_GROUP_NAME, TEST_QUERY_STRING, TEST_START_TIME, TEST_END_TIME)

    assert exc.value.status == result['status']


@pytest.mark.unit
@pytest.mark.parametrize('result', [TEST_RUNNING_QUERY, TEST_UNKNOWN_QUERY])
def test_query_timesout(context, result):
    MockEx(context.mock_aws_service.logs_client.start_query).if_called_with(
        logGroupName=TEST_LOG_GROUP_NAME,
        queryString=TEST_QUERY_STRING,
        startTime=TEST_START_TIME,
        endTime=TEST_END_TIME,
        limit=100,
    ).return_value({'queryId': TEST_QUERY_ID})

    context.mock_aws_service.logs_client.get_query_results.return_value = result

    with pytest.raises(TimeoutError):
        context.cloudwatch.query(TEST_LOG_GROUP_NAME, TEST_QUERY_STRING, TEST_START_TIME, TEST_END_TIME, timeout=2)
