from airflow.operators.docker_plugin import DockerRemovableContainer
from airflow.operators.docker_plugin import DockerWithVariablesOperator
import unittest
from tests.utils.mock_helpers import patch_plugin_file
from datetime import datetime, timedelta
from docker.errors import NotFound, APIError

DAG_ARGS = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2017, 5, 1),
    'catchup_by_default': False,
    'retries': 1,
    'retry_delay': timedelta(seconds=2),
    'retry_exponential_backoff': True,
}
TASK_ID = 'test_docker_with_variables'
IMAGE = 'alpine:latest'
MOUNT_POINT = '/run/secret'
COMMAND_CHECK_MOUNT_TEMPLATE = 'sh -c "mount | grep %s > /dev/null && ls %s"'
COMMAND_SHOW_ITEMS = 'sh -c "ls -1 %s |\
    while read file; do echo $file; cat \\"%s/$file\\"; echo; done"' \
    % (MOUNT_POINT, MOUNT_POINT)


DEFAULT_VARIABLES = {
    'char': 'a',
    'string': 'abcd',
    'key with spaces': '12345',
    'multiline': '''
                multiline
                string''',
}


def variables_to_show_items(variables):
    show_items_builder = []
    for key in sorted(DEFAULT_VARIABLES):
        value = DEFAULT_VARIABLES[key]
        show_items_builder.append(key)
        show_items_builder.append(str(value))
    show_items_builder.append('')
    return '\n'.join(show_items_builder)


class TestDockerConfigurableOperator(unittest.TestCase):
    def test_should_fail_already_removed(self):
        operator = DockerRemovableContainer(
            host_args={'auto_remove': True},
            task_id=TASK_ID,
            default_args=DAG_ARGS,
            image=IMAGE,
            remove=True,
            command='echo fail'
            )
        try:
            with self.assertRaises(APIError):
                operator.execute(None)
        finally:
            try:
                operator.cli.remove_container(operator.container)
            except Exception:
                pass

    def test_should_override_command(self):
        NEW_TEXT = 'NEW_TEXT'
        operator = DockerRemovableContainer(
            container_args={'command': 'echo %s' % NEW_TEXT},
            task_id=TASK_ID,
            default_args=DAG_ARGS,
            image=IMAGE,
            remove=True,
            command='echo old text :(',
            xcom_push=True,
            xcom_all=True,
            )
        assert '%s\n' % NEW_TEXT == operator.execute(None).decode('utf-8')


class TestDockerRemovableContainer(unittest.TestCase):
    def test_should_keep_container(self):
        operator = DockerRemovableContainer(
            task_id=TASK_ID,
            default_args=DAG_ARGS,
            image=IMAGE,
            remove=False,
            command='echo should keep'
            )
        operator.execute(None)

        try:
            operator.cli.inspect_container(operator.container)
        finally:
            try:
                operator.cli.remove_container(operator.container)
            except Exception:
                pass

    def test_should_remove_container_by_default(self):
        operator = DockerRemovableContainer(
            task_id=TASK_ID,
            default_args=DAG_ARGS,
            image=IMAGE,
            command='echo should remove by default'
            )
        operator.execute(None)

        try:
            with self.assertRaises(NotFound):
                operator.cli.inspect_container(operator.container)
        finally:
            try:
                operator.cli.remove_container(operator.container)
            except Exception:
                pass


class TestDockerWithVariables(unittest.TestCase):
    def test_should_mount_and_be_empty_with_default_mount_point(self):
        operator = DockerWithVariablesOperator(
            variables=[],
            task_id=TASK_ID,
            default_args=DAG_ARGS,
            image=IMAGE,
            xcom_push=True,
            xcom_all=True,
            command=COMMAND_CHECK_MOUNT_TEMPLATE % (
                DockerWithVariablesOperator.DEFAULT_MOUNT_POINT,
                DockerWithVariablesOperator.DEFAULT_MOUNT_POINT)
            )
        items = operator.execute(None)  # will fail if no mount found
        assert operator
        assert operator.task_id == TASK_ID
        assert not items  # mount was found but check to make sure it's empty

    def test_should_mount_and_be_empty(self):
        operator = DockerWithVariablesOperator(
            variables=[],
            mount_point=MOUNT_POINT,
            task_id=TASK_ID,
            default_args=DAG_ARGS,
            image=IMAGE,
            xcom_push=True,
            xcom_all=True,
            command=COMMAND_CHECK_MOUNT_TEMPLATE % (MOUNT_POINT, MOUNT_POINT)
            )
        items = operator.execute(None)  # will fail if no mount found
        assert operator
        assert operator.task_id == TASK_ID
        assert not items  # mount was found but check to make sure it's empty

    @patch_plugin_file('plugins/custom/docker', 'Variable', autospec=True)
    def test_should_find_variables(self, variable_class):
        variable_class.get.side_effect = DEFAULT_VARIABLES.__getitem__

        operator = DockerWithVariablesOperator(
            variables=DEFAULT_VARIABLES.keys(),
            mount_point=MOUNT_POINT,
            task_id=TASK_ID,
            default_args=DAG_ARGS,
            image=IMAGE,
            xcom_push=True,
            xcom_all=True,
            command=COMMAND_SHOW_ITEMS
        )

        show_items = operator.execute(None)
        assert operator
        assert operator.task_id == TASK_ID

        assert str(show_items, 'utf-8') == variables_to_show_items(
            DEFAULT_VARIABLES)

    @patch_plugin_file('plugins/custom/docker', 'Variable', autospec=True)
    def test_should_fail_when_variable_not_found(self, variable_class):
        variable_class.get.side_effect = DEFAULT_VARIABLES.__getitem__

        bad_keys = list(DEFAULT_VARIABLES.keys())
        bad_keys.append('bad_key')

        operator = DockerWithVariablesOperator(
            variables=bad_keys,
            mount_point=MOUNT_POINT,
            task_id=TASK_ID,
            default_args=DAG_ARGS,
            image=IMAGE,
            xcom_push=True,
            xcom_all=True,
            command='ls'
        )

        with self.assertRaises(KeyError):
            operator.execute(None)
