#
# 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

from __future__ import absolute_import

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'])
          ]
      },
  ]

  # 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')

  # 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')

  def test_display_data(self):
    for case in PipelineOptionsTest.TEST_CASES:
      options = PipelineOptions(flags=case['flags'])
      dd = DisplayData.create_from(options)