#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#

"""Unit tests for the pipeline options module."""

# pytype: skip-file

import json
import logging
import unittest

import hamcrest as hc

from apache_beam.options.pipeline_options import DebugOptions
from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import ProfilingOptions
from apache_beam.options.pipeline_options import TypeOptions
from apache_beam.options.pipeline_options import WorkerOptions
from apache_beam.options.value_provider import RuntimeValueProvider
from apache_beam.options.value_provider import StaticValueProvider
from apache_beam.transforms.display import DisplayData
from apache_beam.transforms.display_test import DisplayDataItemMatcher


class PipelineOptionsTest(unittest.TestCase):
  def setUp(self):
    # Reset runtime options to avoid side-effects caused by other tests.
    # Note that is_accessible assertions require runtime_options to
    # be uninitialized.
    RuntimeValueProvider.set_runtime_options(None)

  def tearDown(self):
    # Reset runtime options to avoid side-effects in other tests.
    RuntimeValueProvider.set_runtime_options(None)

  TEST_CASES = [
      {
          'flags': ['--num_workers', '5'],
          'expected': {
              'num_workers': 5,
              'mock_flag': False,
              'mock_option': None,
              'mock_multi_option': None
          },
          'display_data': [DisplayDataItemMatcher('num_workers', 5)]
      },
      {
          'flags': ['--direct_num_workers', '5'],
          'expected': {
              'direct_num_workers': 5,
              'mock_flag': False,
              'mock_option': None,
              'mock_multi_option': None
          },
          'display_data': [DisplayDataItemMatcher('direct_num_workers', 5)]
      },
      {
          'flags': ['--direct_running_mode', 'multi_threading'],
          'expected': {
              'direct_running_mode': 'multi_threading',
              'mock_flag': False,
              'mock_option': None,
              'mock_multi_option': None
          },
          'display_data': [
              DisplayDataItemMatcher('direct_running_mode', 'multi_threading')
          ]
      },
      {
          'flags': ['--direct_running_mode', 'multi_processing'],
          'expected': {
              'direct_running_mode': 'multi_processing',
              'mock_flag': False,
              'mock_option': None,
              'mock_multi_option': None
          },
          'display_data': [
              DisplayDataItemMatcher('direct_running_mode', 'multi_processing')
          ]
      },
      {
          'flags': [
              '--profile_cpu', '--profile_location', 'gs://bucket/', 'ignored'
          ],
          'expected': {
              'profile_cpu': True,
              'profile_location': 'gs://bucket/',
              'mock_flag': False,
              'mock_option': None,
              'mock_multi_option': None
          },
          'display_data': [
              DisplayDataItemMatcher('profile_cpu', True),
              DisplayDataItemMatcher('profile_location', 'gs://bucket/')
          ]
      },
      {
          'flags': ['--num_workers', '5', '--mock_flag'],
          'expected': {
              'num_workers': 5,
              'mock_flag': True,
              'mock_option': None,
              'mock_multi_option': None
          },
          'display_data': [
              DisplayDataItemMatcher('num_workers', 5),
              DisplayDataItemMatcher('mock_flag', True)
          ]
      },
      {
          'flags': ['--mock_option', 'abc'],
          'expected': {
              'mock_flag': False,
              'mock_option': 'abc',
              'mock_multi_option': None
          },
          'display_data': [DisplayDataItemMatcher('mock_option', 'abc')]
      },
      {
          'flags': ['--mock_option', ' abc def '],
          'expected': {
              'mock_flag': False,
              'mock_option': ' abc def ',
              'mock_multi_option': None
          },
          'display_data': [DisplayDataItemMatcher('mock_option', ' abc def ')]
      },
      {
          'flags': ['--mock_option= abc xyz '],
          'expected': {
              'mock_flag': False,
              'mock_option': ' abc xyz ',
              'mock_multi_option': None
          },
          'display_data': [DisplayDataItemMatcher('mock_option', ' abc xyz ')]
      },
      {
          'flags': [
              '--mock_option=gs://my bucket/my folder/my file',
              '--mock_multi_option=op1',
              '--mock_multi_option=op2'
          ],
          'expected': {
              'mock_flag': False,
              'mock_option': 'gs://my bucket/my folder/my file',
              'mock_multi_option': ['op1', 'op2']
          },
          'display_data': [
              DisplayDataItemMatcher(
                  'mock_option', 'gs://my bucket/my folder/my file'),
              DisplayDataItemMatcher('mock_multi_option', ['op1', 'op2'])
          ]
      },
      {
          'flags': ['--mock_multi_option=op1', '--mock_multi_option=op2'],
          'expected': {
              'mock_flag': False,
              'mock_option': None,
              'mock_multi_option': ['op1', 'op2']
          },
          'display_data': [
              DisplayDataItemMatcher('mock_multi_option', ['op1', 'op2'])
          ]
      },
      {
          'flags': ['--mock_json_option={"11a": 0, "37a": 1}'],
          'expected': {
              'mock_flag': False,
              'mock_option': None,
              'mock_multi_option': None,
              'mock_json_option': {
                  '11a': 0, '37a': 1
              },
          },
          'display_data': [
              DisplayDataItemMatcher('mock_json_option', {
                  '11a': 0, '37a': 1
              })
          ]
      },
  ]

  # Used for testing newly added flags.
  class MockOptions(PipelineOptions):
    @classmethod
    def _add_argparse_args(cls, parser):
      parser.add_argument('--mock_flag', action='store_true', help='mock flag')
      parser.add_argument('--mock_option', help='mock option')
      parser.add_argument(
          '--mock_multi_option', action='append', help='mock multi option')
      parser.add_argument('--option with space', help='mock option with space')
      parser.add_argument('--mock_json_option', type=json.loads, default={})

  # Use with MockOptions in test cases where multiple option classes are needed.
  class FakeOptions(PipelineOptions):
    @classmethod
    def _add_argparse_args(cls, parser):
      parser.add_argument('--fake_flag', action='store_true', help='fake flag')
      parser.add_argument('--fake_option', help='fake option')
      parser.add_argument(
          '--fake_multi_option', action='append', help='fake multi option')

  @unittest.skip("TODO(BEAM-12515): Flaky test.")
  def test_display_data(self):
    for case in PipelineOptionsTest.TEST_CASES:
      options = PipelineOptions(flags=case['flags'])
      dd = DisplayData.create_from(options)
      hc.assert_that(dd.items, hc.contains_inanyorder(*case['display_data']))

  def test_get_all_options_subclass(self):
    for case in PipelineOptionsTest.TEST_CASES:
      options = PipelineOptionsTest.MockOptions(flags=case['flags'])
      self.assertDictContainsSubset(case['expected'], options.get_all_options())
      self.assertEqual(
          options.view_as(PipelineOptionsTest.MockOptions).mock_flag,
          case['expected']['mock_flag'])
      self.assertEqual(
          options.view_as(PipelineOptionsTest.MockOptions).mock_option,
          case['expected']['mock_option'])
      self.assertEqual(
          options.view_as(PipelineOptionsTest.MockOptions).mock_multi_option,
          case['expected']['mock_multi_option'])

  def test_get_all_options(self):
    for case in PipelineOptionsTest.TEST_CASES:
      options = PipelineOptions(flags=case['flags'])
      self.assertDictContainsSubset(case['expected'], options.get_all_options())
      self.assertEqual(
          options.view_as(PipelineOptionsTest.MockOptions).mock_flag,
          case['expected']['mock_flag'])
      self.assertEqual(
          options.view_as(PipelineOptionsTest.MockOptions).mock_option,
          case['expected']['mock_option'])
      self.assertEqual(
          options.view_as(PipelineOptionsTest.MockOptions).mock_multi_option,
          case['expected']['mock_multi_option'])

  def test_sublcalsses_of_pipeline_options_can_be_instantiated(self):
    for case in PipelineOptionsTest.TEST_CASES:
      mock_options = PipelineOptionsTest.MockOptions(flags=case['flags'])
      self.assertEqual(mock_options.mock_flag, case['expected']['mock_flag'])
      self.assertEqual(
          mock_options.mock_option, case['expected']['mock_option'])
      self.assertEqual(
          mock_options.mock_multi_option, case['expected']['mock_multi_option'])

  def test_views_can_be_constructed_from_pipeline_option_subclasses(self):
    for case in PipelineOptionsTest.TEST_CASES:
      fake_options = PipelineOptionsTest.FakeOptions(flags=case['flags'])
      mock_options = fake_options.view_as(PipelineOptionsTest.MockOptions)

      self.assertEqual(mock_options.mock_flag, case['expected']['mock_flag'])
      self.assertEqual(
          mock_options.mock_option, case['expected']['mock_option'])
      self.assertEqual(
          mock_options.mock_multi_option, case['expected']['mock_multi_option'])

  def test_views_do_not_expose_options_defined_by_other_views(self):
    flags = ['--mock_option=mock_value', '--fake_option=fake_value']

    options = PipelineOptions(flags)
    assert options.view_as(
        PipelineOptionsTest.MockOptions).mock_option == 'mock_value'
    assert options.view_as(
        PipelineOptionsTest.FakeOptions).fake_option == 'fake_value'
    assert options.view_as(PipelineOptionsTest.MockOptions).view_as(
        PipelineOptionsTest.FakeOptions).fake_option == 'fake_value'

    self.assertRaises(
        AttributeError,
        lambda: options.view_as(PipelineOptionsTest.MockOptions).fake_option)
    self.assertRaises(
        AttributeError,
        lambda: options.view_as(PipelineOptionsTest.MockOptions).view_as(
            PipelineOptionsTest.FakeOptions).view_as(
                PipelineOptionsTest.MockOptions).fake_option)

  def test_from_dictionary(self):
    for case in PipelineOptionsTest.TEST_CASES:
      options = PipelineOptions(flags=case['flags'])
      all_options_dict = options.get_all_options()
      options_from_dict = PipelineOptions.from_dictionary(all_options_dict)
      self.assertEqual(
          options_from_dict.view_as(PipelineOptionsTest.MockOptions).mock_flag,
          case['expected']['mock_flag'])
      self.assertEqual(
          options.view_as(PipelineOptionsTest.MockOptions).mock_option,
          case['expected']['mock_option'])
      self.assertEqual(
          options.view_as(PipelineOptionsTest.MockOptions).mock_multi_option,
          case['expected']['mock_multi_option'])
      self.assertEqual(
          options.view_as(PipelineOptionsTest.MockOptions).mock_json_option,
          case['expected'].get('mock_json_option', {}))

  def test_option_with_space(self):
    options = PipelineOptions(flags=['--option with space= value with space'])
    self.assertEqual(
        getattr(
            options.view_as(PipelineOptionsTest.MockOptions),
            'option with space'),
        ' value with space')
    options_from_dict = PipelineOptions.from_dictionary(
        options.get_all_options())
    self.assertEqual(
        getattr(
            options_from_dict.view_as(PipelineOptionsTest.MockOptions),
            'option with space'),
        ' value with space')

  def test_retain_unknown_options_binary_store_string(self):
    options = PipelineOptions(['--unknown_option', 'some_value'])
    result = options.get_all_options(retain_unknown_options=True)
    self.assertEqual(result['unknown_option'], 'some_value')

  def test_retain_unknown_options_binary_equals_store_string(self):
    options = PipelineOptions(['--unknown_option=some_value'])
    result = options.get_all_options(retain_unknown_options=True)
    self.assertEqual(result['unknown_option'], 'some_value')

  def test_retain_unknown_options_binary_multi_equals_store_string(self):
    options = PipelineOptions(['--unknown_option=expr = "2 + 2 = 5"'])
    result = options.get_all_options(retain_unknown_options=True)
    self.assertEqual(result['unknown_option'], 'expr = "2 + 2 = 5"')

  def test_retain_unknown_options_binary_single_dash_store_string(self):
    options = PipelineOptions(['-i', 'some_value'])
    result = options.get_all_options(retain_unknown_options=True)
    self.assertEqual(result['i'], 'some_value')

  def test_retain_unknown_options_unary_store_true(self):
    options = PipelineOptions(['--unknown_option'])
    result = options.get_all_options(retain_unknown_options=True)
    self.assertEqual(result['unknown_option'], True)

  def test_retain_unknown_options_consecutive_unary_store_true(self):
    options = PipelineOptions(['--option_foo', '--option_bar'])
    result = options.get_all_options(retain_unknown_options=True)
    self.assertEqual(result['option_foo'], True)
    self.assertEqual(result['option_bar'], True)

  def test_retain_unknown_options_unary_single_dash_store_true(self):
    options = PipelineOptions(['-i'])
    result = options.get_all_options(retain_unknown_options=True)
    self.assertEqual(result['i'], True)

  def test_retain_unknown_options_unary_missing_prefix(self):
    options = PipelineOptions(['bad_option'])
    with self.assertRaises(SystemExit):
      options.get_all_options(retain_unknown_options=True)

  def test_override_options(self):
    base_flags = ['--num_workers', '5']
    options = PipelineOptions(base_flags)
    self.assertEqual(options.get_all_options()['num_workers'], 5)
    self.assertEqual(options.get_all_options()['mock_flag'], False)

    options.view_as(PipelineOptionsTest.MockOptions).mock_flag = True
    self.assertEqual(options.get_all_options()['num_workers'], 5)
    self.assertTrue(options.get_all_options()['mock_flag'])

  def test_override_init_options(self):
    base_flags = ['--num_workers', '5']
    options = PipelineOptions(base_flags, mock_flag=True)
    self.assertEqual(options.get_all_options()['num_workers'], 5)
    self.assertEqual(options.get_all_options()['mock_flag'], True)

  def test_invalid_override_init_options(self):
    base_flags = ['--num_workers', '5']
    options = PipelineOptions(base_flags, mock_invalid_flag=True)
    self.assertEqual(options.get_all_options()['num_workers'], 5)
    self.assertEqual(options.get_all_options()['mock_flag'], False)

  def test_experiments(self):
    options = PipelineOptions(['--experiment', 'abc', '--experiment', 'def'])
    self.assertEqual(
        sorted(options.get_all_options()['experiments']), ['abc', 'def'])

    options = PipelineOptions(['--experiments', 'abc', '--experiments', 'def'])
    self.assertEqual(
        sorted(options.get_all_options()['experiments']), ['abc', 'def'])

    options = PipelineOptions(flags=[''])
    self.assertEqual(options.get_all_options()['experiments'], None)

  def test_worker_options(self):
    options = PipelineOptions(['--machine_type', 'abc', '--disk_type', 'def'])
    worker_options = options.view_as(WorkerOptions)
    self.assertEqual(worker_options.machine_type, 'abc')
    self.assertEqual(worker_options.disk_type, 'def')

    options = PipelineOptions(
        ['--worker_machine_type', 'abc', '--worker_disk_type', 'def'])
    worker_options = options.view_as(WorkerOptions)
    self.assertEqual(worker_options.machine_type, 'abc')
    self.assertEqual(worker_options.disk_type, 'def')

  def test_option_modifications_are_shared_between_views(self):
    pipeline_options = PipelineOptions([
        '--mock_option',
        'value',
        '--mock_flag',
        '--mock_multi_option',
        'value1',
        '--mock_multi_option',
        'value2',
    ])

    mock_options = PipelineOptionsTest.MockOptions([
        '--mock_option',
        'value',
        '--mock_flag',
        '--mock_multi_option',
        'value1',
        '--mock_multi_option',
        'value2',
    ])

    for options in [pipeline_options, mock_options]:
      view1 = options.view_as(PipelineOptionsTest.MockOptions)
      view2 = options.view_as(PipelineOptionsTest.MockOptions)

      view1.mock_option = 'new_value'
      view1.mock_flag = False
      view1.mock_multi_option.append('value3')

      view3 = options.view_as(PipelineOptionsTest.MockOptions)
      view4 = view1.view_as(PipelineOptionsTest.MockOptions)
      view5 = options.view_as(TypeOptions).view_as(
          PipelineOptionsTest.MockOptions)

      for view in [view1, view2, view3, view4, view5]:
        self.assertEqual('new_value', view.mock_option)
        self.assertFalse(view.mock_flag)
        self.assertEqual(['value1', 'value2', 'value3'], view.mock_multi_option)

  def test_uninitialized_option_modifications_are_shared_between_views(self):
    options = PipelineOptions([])

    view1 = options.view_as(PipelineOptionsTest.MockOptions)
    view2 = options.view_as(PipelineOptionsTest.MockOptions)

    view1.mock_option = 'some_value'
    view1.mock_flag = False
    view1.mock_multi_option = ['value1', 'value2']

    view3 = options.view_as(PipelineOptionsTest.MockOptions)
    view4 = view1.view_as(PipelineOptionsTest.MockOptions)
    view5 = options.view_as(TypeOptions).view_as(
        PipelineOptionsTest.MockOptions)

    for view in [view1, view2, view3, view4, view5]:
      self.assertEqual('some_value', view.mock_option)
      self.assertFalse(view.mock_flag)
      self.assertEqual(['value1', 'value2'], view.mock_multi_option)

  def test_extra_package(self):
    options = PipelineOptions([
        '--extra_package',
        'abc',
        '--extra_packages',
        'def',
        '--extra_packages',
        'ghi'
    ])
    self.assertEqual(
        sorted(options.get_all_options()['extra_packages']),
        ['abc', 'def', 'ghi'])

    options = PipelineOptions(flags=[''])
    self.assertEqual(options.get_all_options()['extra_packages'], None)

  def test_dataflow_job_file(self):
    options = PipelineOptions(['--dataflow_job_file', 'abc'])
    self.assertEqual(options.get_all_options()['dataflow_job_file'], 'abc')

    options = PipelineOptions(flags=[''])
    self.assertEqual(options.get_all_options()['dataflow_job_file'], None)

  def test_template_location(self):
    options = PipelineOptions(['--template_location', 'abc'])
    self.assertEqual(options.get_all_options()['template_location'], 'abc')

    options = PipelineOptions(flags=[''])
    self.assertEqual(options.get_all_options()['template_location'], None)

  def test_redefine_options(self):
    class TestRedefinedOptions(PipelineOptions):  # pylint: disable=unused-variable
      @classmethod
      def _add_argparse_args(cls, parser):
        parser.add_argument('--redefined_flag', action='store_true')

    class TestRedefinedOptions(PipelineOptions):
      @classmethod
      def _add_argparse_args(cls, parser):
        parser.add_argument('--redefined_flag', action='store_true')

    options = PipelineOptions(['--redefined_flag'])
    self.assertTrue(options.get_all_options()['redefined_flag'])

  # TODO(BEAM-1319): Require unique names only within a test.
  # For now, <file name acronym>_vp_arg<number> will be the convention
  # to name value-provider arguments in tests, as opposed to
  # <file name acronym>_non_vp_arg<number> for non-value-provider arguments.
  # The number will grow per file as tests are added.
  def test_value_provider_options(self):
    class UserOptions(PipelineOptions):
      @classmethod
      def _add_argparse_args(cls, parser):
        parser.add_value_provider_argument(
            '--pot_vp_arg1', help='This flag is a value provider')

        parser.add_value_provider_argument('--pot_vp_arg2', default=1, type=int)

        parser.add_argument('--pot_non_vp_arg1', default=1, type=int)

    # Provide values: if not provided, the option becomes of the type runtime vp
    options = UserOptions(['--pot_vp_arg1', 'hello'])
    self.assertIsInstance(options.pot_vp_arg1, StaticValueProvider)
    self.assertIsInstance(options.pot_vp_arg2, RuntimeValueProvider)
    self.assertIsInstance(options.pot_non_vp_arg1, int)

    # Values can be overwritten
    options = UserOptions(
        pot_vp_arg1=5,
        pot_vp_arg2=StaticValueProvider(value_type=str, value='bye'),
        pot_non_vp_arg1=RuntimeValueProvider(
            option_name='foo', value_type=int, default_value=10))
    self.assertEqual(options.pot_vp_arg1, 5)
    self.assertTrue(
        options.pot_vp_arg2.is_accessible(),
        '%s is not accessible' % options.pot_vp_arg2)
    self.assertEqual(options.pot_vp_arg2.get(), 'bye')
    self.assertFalse(options.pot_non_vp_arg1.is_accessible())

    with self.assertRaises(RuntimeError):
      options.pot_non_vp_arg1.get()

  # Converts extra arguments to list value.
  def test_extra_args(self):
    options = PipelineOptions([
        '--extra_arg',
        'val1',
        '--extra_arg',
        'val2',
        '--extra_arg=val3',
        '--unknown_arg',
        'val4'
    ])

    def add_extra_options(parser):
      parser.add_argument("--extra_arg", action='append')

    self.assertEqual(
        options.get_all_options(
            add_extra_args_fn=add_extra_options)['extra_arg'],
        ['val1', 'val2', 'val3'])

  # The argparse package by default tries to autocomplete option names. This
  # results in an "ambiguous option" error from argparse when an unknown option
  # matching multiple known ones are used. This tests that we suppress this
  # error.
  def test_unknown_option_prefix(self):
    # Test that the "ambiguous option" error is suppressed.
    options = PipelineOptions(['--profi', 'val1'])
    options.view_as(ProfilingOptions)

    # Test that valid errors are not suppressed.
    with self.assertRaises(SystemExit):
      # Invalid option choice.
      options = PipelineOptions(['--type_check_strictness', 'blahblah'])
      options.view_as(TypeOptions)

  def test_add_experiment(self):
    options = PipelineOptions([])
    options.view_as(DebugOptions).add_experiment('new_experiment')
    self.assertEqual(['new_experiment'],
                     options.view_as(DebugOptions).experiments)

  def test_add_experiment_preserves_existing_experiments(self):
    options = PipelineOptions(['--experiment=existing_experiment'])
    options.view_as(DebugOptions).add_experiment('new_experiment')
    self.assertEqual(['existing_experiment', 'new_experiment'],
                     options.view_as(DebugOptions).experiments)

  def test_lookup_experiments(self):
    options = PipelineOptions([
        '--experiment=existing_experiment',
        '--experiment',
        'key=value',
        '--experiment',
        'master_key=k1=v1,k2=v2',
    ])
    debug_options = options.view_as(DebugOptions)
    self.assertEqual(
        'default_value',
        debug_options.lookup_experiment('nonexistent', 'default_value'))
    self.assertEqual(
        'value', debug_options.lookup_experiment('key', 'default_value'))
    self.assertEqual(
        'k1=v1,k2=v2', debug_options.lookup_experiment('master_key'))
    self.assertEqual(
        True, debug_options.lookup_experiment('existing_experiment'))

  def test_transform_name_mapping(self):
    options = PipelineOptions(['--transform_name_mapping={\"from\":\"to\"}'])
    mapping = options.view_as(GoogleCloudOptions).transform_name_mapping
    self.assertEqual(mapping['from'], 'to')

  def test_dataflow_service_options(self):
    options = PipelineOptions([
        '--dataflow_service_option',
        'whizz=bang',
        '--dataflow_service_option',
        'beep=boop'
    ])
    self.assertEqual(
        sorted(options.get_all_options()['dataflow_service_options']),
        ['beep=boop', 'whizz=bang'])

    options = PipelineOptions([
        '--dataflow_service_options',
        'whizz=bang',
        '--dataflow_service_options',
        'beep=boop'
    ])
    self.assertEqual(
        sorted(options.get_all_options()['dataflow_service_options']),
        ['beep=boop', 'whizz=bang'])

    options = PipelineOptions(flags=[''])
    self.assertEqual(
        options.get_all_options()['dataflow_service_options'], None)


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  unittest.main()
