# Copyright 2022 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 json
import os
import tempfile
import textwrap
import unittest
from unittest.mock import MagicMock
from unittest.mock import Mock
from unittest.mock import patch

from absl.testing import parameterized
from google.protobuf import json_format
from kfp.client import auth
from kfp.client import client
from kfp.compiler import Compiler
from kfp.dsl import component
from kfp.dsl import pipeline
from kfp.pipeline_spec import pipeline_spec_pb2
import kfp_server_api
import yaml


class TestValidatePipelineName(parameterized.TestCase):

    @parameterized.parameters([
        'pipeline',
        'my-pipeline',
        'my-pipeline-1',
        '1pipeline',
        'pipeline1',
        'my_pipeline',
        "person's-pipeline",
        'my pipeline',
        'pipeline.yaml',
    ])
    def test_valid(self, name: str):
        client.validate_pipeline_display_name(name)

    @parameterized.parameters(['', '   ', '\t'])
    def test_invalid(self, name: str):
        with self.assertRaisesRegex(
                ValueError,
                'Invalid pipeline name. Pipeline name cannot be empty or contain only whitespace.'
        ):
            client.validate_pipeline_display_name(name)


class TestOverrideCachingOptions(parameterized.TestCase):

    def test_override_caching_of_multiple_components(self):

        @component
        def hello_word(text: str) -> str:
            return text

        @component
        def to_lower(text: str) -> str:
            return text.lower()

        @pipeline(
            name='sample two-step pipeline',
            description='a minimal two-step pipeline')
        def pipeline_with_two_component(text: str = 'hi there'):

            component_1 = hello_word(text=text).set_caching_options(True)
            component_2 = to_lower(
                text=component_1.output).set_caching_options(False)

        with tempfile.TemporaryDirectory() as tempdir:
            temp_filepath = os.path.join(tempdir, 'hello_world_pipeline.yaml')
            Compiler().compile(
                pipeline_func=pipeline_with_two_component,
                package_path=temp_filepath)

            with open(temp_filepath, 'r') as f:
                pipeline_obj = yaml.safe_load(f)
                pipeline_spec = json_format.ParseDict(
                    pipeline_obj, pipeline_spec_pb2.PipelineSpec())
                client._override_caching_options(pipeline_spec, True)
                pipeline_obj = json_format.MessageToDict(pipeline_spec)
                self.assertTrue(pipeline_obj['root']['dag']['tasks']
                                ['hello-word']['cachingOptions']['enableCache'])
                self.assertTrue(pipeline_obj['root']['dag']['tasks']['to-lower']
                                ['cachingOptions']['enableCache'])
                self.assertFalse('cacheKey' in pipeline_obj['root']['dag']
                                 ['tasks']['hello-word']['cachingOptions'])
                self.assertFalse('cacheKey' in pipeline_obj['root']['dag']
                                 ['tasks']['to-lower']['cachingOptions'])

    def test_override_cache_key(self):

        @component
        def hello_word(text: str) -> str:
            return text

        @component
        def to_lower(text: str) -> str:
            return text.lower()

        @pipeline(
            name='sample two-step pipeline',
            description='a minimal two-step pipeline')
        def pipeline_with_two_component(text: str = 'hi there'):
            component_1 = hello_word(text=text).set_caching_options(True)
            component_2 = to_lower(
                text=component_1.output).set_caching_options(False)

        with tempfile.TemporaryDirectory() as tempdir:
            temp_filepath = os.path.join(tempdir, 'hello_world_pipeline.yaml')
            Compiler().compile(
                pipeline_func=pipeline_with_two_component,
                package_path=temp_filepath)

            with open(temp_filepath, 'r') as f:
                pipeline_obj = yaml.safe_load(f)
                pipeline_spec = json_format.ParseDict(
                    pipeline_obj, pipeline_spec_pb2.PipelineSpec())
                client._override_caching_options(
                    pipeline_spec, True, cache_key='OVERRIDE_KEY')
                pipeline_obj = json_format.MessageToDict(pipeline_spec)
                self.assertTrue(pipeline_obj['root']['dag']['tasks']
                                ['hello-word']['cachingOptions']['enableCache'])
                self.assertEqual(
                    pipeline_obj['root']['dag']['tasks']['hello-word']
                    ['cachingOptions']['cacheKey'], 'OVERRIDE_KEY')
                self.assertTrue(pipeline_obj['root']['dag']['tasks']['to-lower']
                                ['cachingOptions']['enableCache'])
                self.assertEqual(
                    pipeline_obj['root']['dag']['tasks']['to-lower']
                    ['cachingOptions']['cacheKey'], 'OVERRIDE_KEY')


class TestExtractPipelineYAML(parameterized.TestCase):

    def test_extract_pipeline_yaml_single_doc(self):

        with tempfile.TemporaryDirectory() as tempdir:
            temp_filepath = os.path.join(tempdir, 'single_doc_pipeline.yaml')
            with open(temp_filepath, 'w') as f:
                f.write(
                    textwrap.dedent('''
                        components:
                          comp-foo:
                            executorLabel: exec-foo
                        deploymentSpec:
                          executors:
                            exec-foo:
                              container:
                                command:
                                - sh
                                - -c
                                - cat /data/file.txt
                                image: alpine
                        pipelineInfo:
                          name: my-pipeline
                        root:
                          dag:
                            tasks:
                              foo:
                                componentRef:
                                  name: comp-foo
                                taskInfo:
                                  name: foo
                        schemaVersion: 2.1.0
                        sdkVersion: kfp-2.0.0-beta.13
                        '''))

            pipeline_dict = client._extract_pipeline_yaml(
                temp_filepath).to_dict()
            self.assertEqual('my-pipeline',
                             pipeline_dict['pipelineInfo']['name'])

    def test_extract_pipeline_yaml_multiple_docs(self):

        with tempfile.TemporaryDirectory() as tempdir:
            temp_filepath = os.path.join(tempdir, 'multi_docs_pipeline.yaml')
            with open(temp_filepath, 'w') as f:
                f.write(
                    textwrap.dedent('''
                        components:
                          comp-foo:
                            executorLabel: exec-foo
                        deploymentSpec:
                          executors:
                            exec-foo:
                              container:
                                command:
                                - sh
                                - -c
                                - cat /data/file.txt
                                image: alpine
                        pipelineInfo:
                          name: my-pipeline
                        root:
                          dag:
                            tasks:
                              foo:
                                componentRef:
                                  name: comp-foo
                                taskInfo:
                                  name: foo
                        schemaVersion: 2.1.0
                        sdkVersion: kfp-2.0.0-beta.13
                        ---
                        platforms:
                          kubernetes:
                            deploymentSpec:
                              executors:
                                exec-foo:
                                  pvcMount:
                                  - mountPath: /data
                                    constant: my-pvc
                        '''))

            pipeline_dict = client._extract_pipeline_yaml(
                temp_filepath).to_dict()
            self.assertEqual(
                'my-pipeline',
                pipeline_dict['pipeline_spec']['pipelineInfo']['name'])
            self.assertEqual(
                'my-pvc', pipeline_dict['platform_spec']['platforms']
                ['kubernetes']['deploymentSpec']['executors']['exec-foo']
                ['pvcMount'][0]['constant'])


class TestClient(parameterized.TestCase):

    def setUp(self):
        self.client = client.Client(namespace='ns1')

    def test_wait_for_run_completion_invalid_token_should_raise_error(self):
        with self.assertRaises(kfp_server_api.ApiException):
            with patch.object(
                    self.client._run_api,
                    'run_service_get_run',
                    side_effect=kfp_server_api.ApiException) as mock_get_run:
                self.client.wait_for_run_completion(
                    run_id='foo', timeout=1, sleep_duration=0)
                mock_get_run.assert_called_once()

    def test_wait_for_run_completion_expired_access_token(self):
        with patch.object(self.client._run_api,
                          'run_service_get_run') as mock_get_run:
            # We need to iterate through multiple side effects in order to test this logic.
            mock_get_run.side_effect = [
                Mock(state='unknown state'),
                kfp_server_api.ApiException(status=401),
                Mock(state='succeeded'),
            ]

            with patch.object(self.client, '_refresh_api_client_token'
                             ) as mock_refresh_api_client_token:
                self.client.wait_for_run_completion(
                    run_id='foo', timeout=1, sleep_duration=0)
                mock_get_run.assert_called_with(run_id='foo')
                mock_refresh_api_client_token.assert_called_once()

    def test_wait_for_run_completion_valid_token(self):
        with patch.object(self.client._run_api,
                          'run_service_get_run') as mock_get_run:
            mock_get_run.return_value = Mock(state='succeeded')
            response = self.client.wait_for_run_completion(
                run_id='foo', timeout=1, sleep_duration=0)
            mock_get_run.assert_called_once_with(run_id='foo')
            assert response == mock_get_run.return_value

    def test_wait_for_run_completion_run_timeout_should_raise_error(self):
        with self.assertRaises(TimeoutError):
            with patch.object(self.client._run_api,
                              'run_service_get_run') as mock_get_run:
                mock_get_run.return_value = Mock(run=Mock(status='foo'))
                self.client.wait_for_run_completion(
                    run_id='foo', timeout=1, sleep_duration=0)
                mock_get_run.assert_called_once_with(run_id='foo')

    @patch('kfp.Client.get_experiment', side_effect=ValueError)
    def test_create_experiment_no_experiment_should_raise_error(
            self, mock_get_experiment):
        with self.assertRaises(ValueError):
            self.client.create_experiment(name='foo', namespace='ns1')
            mock_get_experiment.assert_called_once_with(
                name='foo', onamespace='ns1')

    @patch('kfp.Client.get_experiment', return_value=Mock(id='foo'))
    @patch('kfp.Client._get_url_prefix', return_value='/pipeline')
    def test_create_experiment_existing_experiment(self, mock_get_url_prefix,
                                                   mock_get_experiment):
        self.client.create_experiment(name='foo')
        mock_get_experiment.assert_called_once_with(
            experiment_name='foo', namespace='ns1')
        mock_get_url_prefix.assert_called_once()

    @patch('kfp_server_api.V2beta1Experiment')
    @patch(
        'kfp.Client.get_experiment',
        side_effect=ValueError('No experiment is found with name'))
    @patch('kfp.Client._get_url_prefix', return_value='/pipeline')
    def test__create_experiment_name_not_found(self, mock_get_url_prefix,
                                               mock_get_experiment,
                                               mock_api_experiment):
        # experiment with the specified name is not found, so a new experiment
        # is created.
        with patch.object(
                self.client._experiment_api,
                'experiment_service_create_experiment',
                return_value=Mock(
                    experiment_id='foo')) as mock_create_experiment:
            self.client.create_experiment(name='foo')
            mock_get_experiment.assert_called_once_with(
                experiment_name='foo', namespace='ns1')
            mock_api_experiment.assert_called_once()
            mock_create_experiment.assert_called_once()
            mock_get_url_prefix.assert_called_once()

    def test_get_experiment_no_experiment_id_or_name_should_raise_error(self):
        with self.assertRaises(ValueError):
            self.client.get_experiment()

    @patch('kfp.Client.get_user_namespace', return_value=None)
    def test_get_experiment_does_not_exist_should_raise_error(
            self, mock_get_user_namespace):
        with self.assertRaises(ValueError):
            with patch.object(
                    self.client._experiment_api,
                    'experiment_service_list_experiments',
                    return_value=Mock(
                        experiments=None)) as mock_list_experiments:
                self.client.get_experiment(experiment_name='foo')
                mock_list_experiments.assert_called_once()
                mock_get_user_namespace.assert_called_once()

    @patch('kfp.Client.get_user_namespace', return_value=None)
    def test_get_experiment_multiple_experiments_with_name_should_raise_error(
            self, mock_get_user_namespace):
        with self.assertRaises(ValueError):
            with patch.object(
                    self.client._experiment_api,
                    'experiment_service_list_experiments',
                    return_value=Mock(
                        experiments=['foo', 'foo'])) as mock_list_experiments:
                self.client.get_experiment(experiment_name='foo')
                mock_list_experiments.assert_called_once()
                mock_get_user_namespace.assert_called_once()

    def test_get_experiment_with_experiment_id(self):
        with patch.object(
                self.client._experiment_api,
                'experiment_service_get_experiment') as mock_get_experiment:
            self.client.get_experiment(experiment_id='foo')
            mock_get_experiment.assert_called_once_with(experiment_id='foo')

    def test_get_experiment_with_experiment_name_and_namespace(self):
        with patch.object(
                self.client._experiment_api,
                'experiment_service_list_experiments') as mock_list_experiments:
            self.client.get_experiment(experiment_name='foo', namespace='ns1')
            mock_list_experiments.assert_called_once()

    @patch('kfp.Client.get_user_namespace', return_value=None)
    def test_get_experiment_with_experiment_name_and_no_namespace(
            self, mock_get_user_namespace):
        with patch.object(
                self.client._experiment_api,
                'experiment_service_list_experiments') as mock_list_experiments:
            self.client.get_experiment(experiment_name='foo')
            mock_list_experiments.assert_called_once()
            mock_get_user_namespace.assert_called_once()

    @patch('kfp_server_api.HealthzServiceApi.healthz_service_get_healthz')
    def test_get_kfp_healthz(self, mock_get_kfp_healthz):
        mock_get_kfp_healthz.return_value = json.dumps([{'foo': 'bar'}])
        response = self.client.get_kfp_healthz()
        mock_get_kfp_healthz.assert_called_once()
        assert (response == mock_get_kfp_healthz.return_value)

    @patch(
        'kfp_server_api.HealthzServiceApi.healthz_service_get_healthz',
        side_effect=kfp_server_api.ApiException)
    def test_get_kfp_healthz_should_raise_error(self, mock_get_kfp_healthz):
        with self.assertRaises(TimeoutError):
            self.client.get_kfp_healthz(sleep_duration=0)
            mock_get_kfp_healthz.assert_called()

    def test_upload_pipeline_without_name(self):

        @component
        def return_bool(boolean: bool) -> bool:
            return boolean

        @pipeline(name='test-upload-without-name', description='description')
        def pipeline_test_upload_without_name(boolean: bool = True):
            return_bool(boolean=boolean)

        with patch.object(self.client._upload_api,
                          'upload_pipeline') as mock_upload_pipeline:
            with patch.object(auth, 'is_ipython', return_value=False):
                with tempfile.TemporaryDirectory() as tmp_path:
                    pipeline_test_path = os.path.join(tmp_path, 'test.yaml')
                    Compiler().compile(
                        pipeline_func=pipeline_test_upload_without_name,
                        package_path=pipeline_test_path)
                    self.client.upload_pipeline(
                        pipeline_package_path=pipeline_test_path,
                        description='description',
                        namespace='ns1')
                    mock_upload_pipeline.assert_called_once_with(
                        pipeline_test_path,
                        name='test-upload-without-name',
                        description='description',
                        namespace='ns1')

    def test_upload_pipeline_from_pipeline_func(self):

        @component
        def return_bool(boolean: bool) -> bool:
            return boolean

        @pipeline(name='test-upload', description='description')
        def pipeline_test_upload_without_name(boolean: bool = True):
            return_bool(boolean=boolean)

        fake_tmpdir = '/fake/tmp'
        fake_pipeline_path = os.path.join(fake_tmpdir, 'pipeline.yaml')
        expected_result = MagicMock()

        with patch('tempfile.TemporaryDirectory') as mock_tempdir, \
             patch('kfp.compiler.Compiler.compile') as mock_compile, \
             patch.object(self.client, 'upload_pipeline', return_value=expected_result) as mock_upload:

            mock_tempdir.return_value.__enter__.return_value = fake_tmpdir

            result = self.client.upload_pipeline_from_pipeline_func(
                pipeline_func=pipeline_test_upload_without_name,
                pipeline_name='test-pipeline',
                description='A test pipeline',
                namespace='test-ns')

            mock_compile.assert_called_once_with(
                pipeline_func=pipeline_test_upload_without_name,
                package_path=fake_pipeline_path,
            )

            mock_upload.assert_called_once_with(
                pipeline_package_path=fake_pipeline_path,
                pipeline_name='test-pipeline',
                description='A test pipeline',
                namespace='test-ns')

            self.assertEqual(result, expected_result)

    @parameterized.parameters([
        'pipeline',
        'my-pipeline',
        'my-pipeline-1',
        '1pipeline',
        'pipeline1',
        'my_pipeline',
        "person's-pipeline",
        'my pipeline',
        'pipeline.yaml',
    ])
    def test_upload_pipeline_with_name(self, pipeline_name):
        with patch.object(self.client._upload_api,
                          'upload_pipeline') as mock_upload_pipeline:
            with patch.object(auth, 'is_ipython', return_value=False):
                self.client.upload_pipeline(
                    pipeline_package_path='fake.yaml',
                    pipeline_name=pipeline_name,
                    description='description',
                    namespace='ns1')
                mock_upload_pipeline.assert_called_once_with(
                    'fake.yaml',
                    name=pipeline_name,
                    description='description',
                    namespace='ns1')

    @parameterized.parameters([
        '',
        '   ',
        '\t',
    ])
    def test_upload_pipeline_with_name_invalid(self, pipeline_name):
        with patch.object(self.client._upload_api,
                          'upload_pipeline') as mock_upload_pipeline:
            with patch.object(auth, 'is_ipython', return_value=False):
                with self.assertRaisesRegex(
                        ValueError,
                        'Invalid pipeline name. Pipeline name cannot be empty or contain only whitespace.'
                ):
                    self.client.upload_pipeline(
                        pipeline_package_path='fake.yaml',
                        pipeline_name=pipeline_name,
                        description='description',
                        namespace='ns1')

    def test_upload_pipeline_version_from_pipeline_func(self):

        @component
        def return_bool(boolean: bool) -> bool:
            return boolean

        @pipeline(name='test-upload', description='description')
        def pipeline_test_upload_without_name(boolean: bool = True):
            return_bool(boolean=boolean)

        fake_tmpdir = '/mock/tmp'
        fake_pipeline_path = os.path.join(fake_tmpdir, 'pipeline.yaml')
        expected_result = MagicMock()

        with patch('tempfile.TemporaryDirectory') as mock_tempdir, \
             patch('kfp.compiler.Compiler.compile') as mock_compile, \
             patch.object(self.client, 'upload_pipeline_version', return_value=expected_result) as mock_upload_version:

            mock_tempdir.return_value.__enter__.return_value = fake_tmpdir

            result = self.client.upload_pipeline_version_from_pipeline_func(
                pipeline_func=pipeline_test_upload_without_name,
                pipeline_version_name='v1.0',
                pipeline_id='pipeline123',
                pipeline_name='test-pipeline',
                description='A test pipeline version')

            mock_compile.assert_called_once_with(
                pipeline_func=pipeline_test_upload_without_name,
                package_path=fake_pipeline_path,
            )

            mock_upload_version.assert_called_once_with(
                pipeline_package_path=fake_pipeline_path,
                pipeline_version_name='v1.0',
                pipeline_id='pipeline123',
                pipeline_name='test-pipeline',
                description='A test pipeline version')

            self.assertEqual(result, expected_result)


if __name__ == '__main__':
    unittest.main()
