"""Tests for data sanitisation."""
from io import StringIO
import json

from botocore.response import StreamingBody
from botocore.stub import Stubber
import pytest

from common.connectors.boto_clients import lambda_client
from common.logic.sanitiser import sanitise
from common.logic.sanitiser import DataSanitisationFailed


@pytest.fixture
def lambda_stubber():
    """Return stubber for the RDS client."""
    with Stubber(lambda_client) as stubber:
        yield stubber


def test_sanitise_data_success(lambda_stubber):
    """Test successful data sanitisation."""
    response_payload = json.dumps({
        'scripts_run': 2,
        'users_created': 5
    })
    lambda_stubber.add_response('invoke', {
        'Payload': StreamingBody(
            StringIO(response_payload),
            len(response_payload)
        )
    }, {
        'FunctionName': 'sanitise-data',
        'InvocationType': 'RequestResponse',
        'Payload': json.dumps({
            'scripts_dir': 'scripts',
            'engine': 'aurora-mysql',
            'users_source_db': {
                'name': 'src-db',
                'type': 'cluster'
            },
            'target_db': {
                'name': 'target-db',
                'type': 'cluster'
            }
        })
    })

    result = sanitise(
        'sanitise-data', 'src-db', 'target-db', 'cluster', 'scripts',
        'aurora-mysql')

    assert result == {
        'scripts_run': 2,
        'users_created': 5
    }
    lambda_stubber.assert_no_pending_responses()


def test_sanitise_data_failure(lambda_stubber):
    """Test unsuccessful data sanitisation."""
    lambda_stubber.add_response('invoke', {
        'FunctionError': 'Sanitisation failed'
    }, {
        'FunctionName': 'sanitise-data',
        'InvocationType': 'RequestResponse',
        'Payload': json.dumps({
            'scripts_dir': 'scripts',
            'engine': 'aurora-mysql',
            'users_source_db': {
                'name': 'src-db',
                'type': 'cluster'
            },
            'target_db': {
                'name': 'target-db',
                'type': 'cluster'
            }
        })
    })

    with pytest.raises(DataSanitisationFailed):
        sanitise('sanitise-data', 'src-db', 'target-db', 'cluster', 'scripts',
                 'aurora-mysql')


def test_sanitise_data_retry_on_inactive_function(lambda_stubber):
    """Test retry when function is inactive."""
    lambda_stubber.add_client_error(
        'invoke', 'CodeArtifactUserPendingException')
    lambda_stubber.add_response('get_function', {
        'Configuration': {
            'State': 'Inactive'
        }
    })
    lambda_stubber.add_response('get_function', {
        'Configuration': {
            'State': 'Active'
        }
    })
    response_payload = json.dumps({
        'scripts_run': 2,
        'users_created': 5
    })
    lambda_stubber.add_response('invoke', {
        'Payload': StreamingBody(
            StringIO(response_payload),
            len(response_payload)
        )
    })

    result = sanitise(
        'sanitise-data', 'src-db', 'target-db', 'cluster', 'scripts',
        'aurora-mysql')

    assert result == {
        'scripts_run': 2,
        'users_created': 5
    }
    lambda_stubber.assert_no_pending_responses()