# Copyright 2018 The Kubeflow Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import unittest

from kfp.deprecated.containers._component_builder import \
    _dependency_to_requirements
from kfp.deprecated.containers._component_builder import _generate_dockerfile
from kfp.deprecated.containers._component_builder import DependencyHelper
from kfp.deprecated.containers._component_builder import VersionedDependency


class TestVersionedDependency(unittest.TestCase):

    def test_version(self):
        """test version overrides min_version and max_version."""
        version = VersionedDependency(
            name='tensorflow',
            version='0.3.0',
            min_version='0.1.0',
            max_version='0.4.0')
        self.assertTrue(version.min_version == '0.3.0')
        self.assertTrue(version.max_version == '0.3.0')
        self.assertTrue(version.has_versions())
        self.assertTrue(version.name == 'tensorflow')

    def test_minmax_version(self):
        """test if min_version and max_version are configured when version is
        not given."""
        version = VersionedDependency(
            name='tensorflow', min_version='0.1.0', max_version='0.4.0')
        self.assertTrue(version.min_version == '0.1.0')
        self.assertTrue(version.max_version == '0.4.0')
        self.assertTrue(version.has_versions())

    def test_min_or_max_version(self):
        """test if min_version and max_version are configured when version is
        not given."""
        version = VersionedDependency(name='tensorflow', min_version='0.1.0')
        self.assertTrue(version.min_version == '0.1.0')
        self.assertTrue(version.has_versions())
        version = VersionedDependency(name='tensorflow', max_version='0.3.0')
        self.assertTrue(version.max_version == '0.3.0')
        self.assertTrue(version.has_versions())

    def test_no_version(self):
        """test the no version scenario."""
        version = VersionedDependency(name='tensorflow')
        self.assertFalse(version.has_min_version())
        self.assertFalse(version.has_max_version())
        self.assertFalse(version.has_versions())


class TestDependencyHelper(unittest.TestCase):

    def test_generate_requirement(self):
        """Test generating requirement file."""

        # prepare
        test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')
        temp_file = os.path.join(test_data_dir, 'test_requirements.tmp')

        dependency_helper = DependencyHelper()
        dependency_helper.add_python_package(
            dependency=VersionedDependency(
                name='tensorflow', min_version='0.10.0', max_version='0.11.0'))
        dependency_helper.add_python_package(
            dependency=VersionedDependency(
                name='kubernetes', min_version='0.6.0'))
        dependency_helper.add_python_package(
            dependency=VersionedDependency(name='pytorch', max_version='0.3.0'))
        dependency_helper.generate_pip_requirements(temp_file)

        golden_requirement_payload = '''\
tensorflow >= 0.10.0, <= 0.11.0
kubernetes >= 0.6.0
pytorch <= 0.3.0
'''
        with open(temp_file, 'r') as f:
            target_requirement_payload = f.read()
        self.assertEqual(target_requirement_payload, golden_requirement_payload)
        os.remove(temp_file)

    def test_add_python_package(self):
        """Test add_python_package."""

        # prepare
        test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')
        temp_file = os.path.join(test_data_dir, 'test_requirements.tmp')

        dependency_helper = DependencyHelper()
        dependency_helper.add_python_package(
            dependency=VersionedDependency(
                name='tensorflow', min_version='0.10.0', max_version='0.11.0'))
        dependency_helper.add_python_package(
            dependency=VersionedDependency(
                name='kubernetes', min_version='0.6.0'))
        dependency_helper.add_python_package(
            dependency=VersionedDependency(
                name='tensorflow', min_version='0.12.0'),
            override=True)
        dependency_helper.add_python_package(
            dependency=VersionedDependency(
                name='kubernetes', min_version='0.8.0'),
            override=False)
        dependency_helper.add_python_package(
            dependency=VersionedDependency(name='pytorch', version='0.3.0'))
        dependency_helper.generate_pip_requirements(temp_file)
        golden_requirement_payload = '''\
tensorflow >= 0.12.0
kubernetes >= 0.6.0
pytorch >= 0.3.0, <= 0.3.0
'''
        with open(temp_file, 'r') as f:
            target_requirement_payload = f.read()
        self.assertEqual(target_requirement_payload, golden_requirement_payload)
        os.remove(temp_file)


class TestGenerator(unittest.TestCase):

    def test_generate_dockerfile(self):
        """Test generate dockerfile."""
        # prepare
        test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')
        target_dockerfile = os.path.join(test_data_dir,
                                         'component.temp.dockerfile')
        golden_dockerfile_payload_one = '''\
FROM gcr.io/ngao-mlpipeline-testing/tensorflow:1.10.0
RUN apt-get update -y && apt-get install --no-install-recommends -y -q python3 python3-pip python3-setuptools
ADD main.py /ml/main.py
'''
        golden_dockerfile_payload_two = '''\
FROM gcr.io/ngao-mlpipeline-testing/tensorflow:1.10.0
RUN apt-get update -y && apt-get install --no-install-recommends -y -q python3 python3-pip python3-setuptools
ADD requirements.txt /ml/requirements.txt
RUN python3 -m pip install -r /ml/requirements.txt
ADD main.py /ml/main.py
'''

        # check
        _generate_dockerfile(
            filename=target_dockerfile,
            base_image='gcr.io/ngao-mlpipeline-testing/tensorflow:1.10.0',
            add_files={'main.py': '/ml/main.py'})
        with open(target_dockerfile, 'r') as f:
            target_dockerfile_payload = f.read()
        self.assertEqual(target_dockerfile_payload,
                         golden_dockerfile_payload_one)

        _generate_dockerfile(
            filename=target_dockerfile,
            base_image='gcr.io/ngao-mlpipeline-testing/tensorflow:1.10.0',
            requirement_filename='requirements.txt',
            add_files={'main.py': '/ml/main.py'})
        with open(target_dockerfile, 'r') as f:
            target_dockerfile_payload = f.read()
        self.assertEqual(target_dockerfile_payload,
                         golden_dockerfile_payload_two)

        # clean up
        os.remove(target_dockerfile)

    def test_generate_requirement(self):
        # prepare
        test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')
        temp_file = os.path.join(test_data_dir, 'test_requirements.tmp')

        dependencies = [
            VersionedDependency(
                name='tensorflow', min_version='0.10.0', max_version='0.11.0'),
            VersionedDependency(name='kubernetes', min_version='0.6.0'),
        ]
        _dependency_to_requirements(dependencies, filename=temp_file)
        golden_payload = '''\
tensorflow >= 0.10.0, <= 0.11.0
kubernetes >= 0.6.0
'''
        with open(temp_file, 'r') as f:
            target_payload = f.read()
        self.assertEqual(target_payload, golden_payload)
        os.remove(temp_file)