text: # Copyright 2014 PerfKitBenchmarker Authors. All rights reserved.
  #
  # 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.
  
  """Tests for beam_benchmark_helper."""
  
  import mock
  import tempfile
  import unittest
  
  from perfkitbenchmarker import beam_benchmark_helper
  from perfkitbenchmarker import dpb_service
  from perfkitbenchmarker import flags
  from perfkitbenchmarker import vm_util
  
  FLAGS = flags.FLAGS
  
  
  class BeamBenchmarkHelperTestCase(unittest.TestCase):
  
    def setUp(self):
      # TODO(ferneyhough): See exactly why this is needed and find a better way
      # to do this. Unittests in PKB should not have to add this call manually.
      FLAGS.mark_as_parsed()
  
    def test_runner_option_override_use_override(self):
      testOptionVal = "TestVal"
      actual_options = []
      beam_benchmark_helper.AddRunnerPipelineOption(actual_options, None, testOptionVal)
      self.assertListEqual(["--runner=" + testOptionVal], actual_options)
  
    def test_runner_option_override_empty_override(self):
      testOptionVal = ""
      actual_options = []
      beam_benchmark_helper.AddRunnerPipelineOption(actual_options, None, testOptionVal)
      self.assertListEqual([], actual_options)
  
    def test_dataflow_runner_name_added(self):
      testOptionVal = "dataflow"
      actual_command = []
      beam_benchmark_helper.AddRunnerArgument(actual_command, testOptionVal)
      self.assertListEqual(["-DintegrationTestRunner=" + testOptionVal], actual_command)
  
    def test_direct_runner_name_added(self):
      testOptionVal = "direct"
      actual_command = []
      beam_benchmark_helper.AddRunnerArgument(actual_command, testOptionVal)
      self.assertListEqual(["-DintegrationTestRunner=" + testOptionVal], actual_command)
  
  
    def test_runner_name_empty(self):
      testOptionVal = ""
      actual_command = []
      beam_benchmark_helper.AddRunnerArgument(actual_command, testOptionVal)
      self.assertListEqual([], actual_command)
  
    def test_extra_property_empty_property(self):
      testOptionVal = ""
      actual_command = []
      beam_benchmark_helper.AddExtraProperties(actual_command, testOptionVal)
      self.assertListEqual([], actual_command)
  
    def test_extra_property_single_property(self):
      testOptionVal = "[key=value]"
      actual_mvn_command = []
      beam_benchmark_helper.AddExtraProperties(actual_mvn_command, testOptionVal)
      self.assertListEqual(["-Dkey=value"], actual_mvn_command)
  
    def test_extra_property_single_property_quoted(self):
      testOptionVal = "[\"key=value\"]"
      actual_mvn_command = []
      beam_benchmark_helper.AddExtraProperties(actual_mvn_command, testOptionVal)
      self.assertListEqual(["-Dkey=value"], actual_mvn_command)
  
    def test_extra_property_multiple_properties(self):
      testOptionVal = "[\"key=value\", \"key2=value2\"]"
      actual_mvn_command = []
      beam_benchmark_helper.AddExtraProperties(actual_mvn_command, testOptionVal)
      self.assertListEqual(["-Dkey=value", "-Dkey2=value2"], actual_mvn_command)
  
    def test_integrationPipelineOptions_rejection(self):
      testOptionVal = "[\"integrationTestPipelineOptions=...\"]"
      actual_mvn_command = []
      with self.assertRaises(ValueError):
        beam_benchmark_helper.AddExtraProperties(actual_mvn_command, testOptionVal)
  
    def test_hdfs_filesystem_addition(self):
      testOptionVal = "hdfs"
      actual_command = []
      beam_benchmark_helper.AddFilesystemArgument(actual_command, testOptionVal)
      self.assertListEqual(["-Dfilesystem=hdfs"], actual_command)
  
    def test_empty_filesystem(self):
      testOptionVal = ""
      actual_command = []
      beam_benchmark_helper.AddFilesystemArgument(actual_command, testOptionVal)
      self.assertListEqual([], actual_command)
  
    def test_add_module(self):
      testOptionVal = "sdks/java/io/tests"
      actual_command = []
      beam_benchmark_helper.AddModuleArgument(actual_command, testOptionVal)
      self.assertListEqual(['-p', testOptionVal], actual_command)
  
    def test_add_empty_module(self):
      testOptionVal = ""
      actual_command = []
      beam_benchmark_helper.AddModuleArgument(actual_command, testOptionVal)
      self.assertListEqual([], actual_command)
  
    def test_initialize_beam_repo_beam_exists(self):
      FLAGS.beam_location = tempfile.mkdtemp()
  
      with mock.patch.object(beam_benchmark_helper, '_PrebuildBeam') as mock_prebuild, \
          mock.patch.object(vm_util, 'GenTempDir'):
  
        mock_spec = mock.MagicMock()
        mock_spec.dpb_service.SERVICE_TYPE = dpb_service.DATAFLOW
  
        beam_benchmark_helper.InitializeBeamRepo(mock_spec)
        mock_prebuild.assert_called_once()
  
    def test_initialize_beam_repo_beam_not_exists(self):
      FLAGS.beam_location = None
  
      with mock.patch.object(beam_benchmark_helper, '_PrebuildBeam') as mock_prebuild, \
          mock.patch.object(vm_util, 'GenTempDir'), \
          mock.patch.object(vm_util, 'GetTempDir'), \
          mock.patch.object(vm_util, 'IssueCommand') as mock_run:
  
        mock_spec = mock.MagicMock()
        mock_spec.dpb_service.SERVICE_TYPE = dpb_service.DATAFLOW
  
        beam_benchmark_helper.InitializeBeamRepo(mock_spec)
  
        expected_cmd = ['git', 'clone', 'https://github.com/apache/beam.git']
        mock_run.assert_called_once_with(expected_cmd,
                                         cwd=vm_util.GetTempDir())
        mock_prebuild.assert_called_once()
  
    def test_beam_prebuild(self):
      FLAGS.beam_prebuilt = False
      FLAGS.beam_it_module = 'sdks/java'
      FLAGS.beam_runner = 'dataflow'
      FLAGS.beam_filesystem = 'hdfs'
      FLAGS.beam_extra_properties = '[extra_key=extra_value]'
  
      with mock.patch.object(beam_benchmark_helper, '_GetGradleCommand') as mock_gradle, \
          mock.patch.object(beam_benchmark_helper, '_GetBeamDir'), \
          mock.patch.object(vm_util, 'IssueCommand') as mock_run:
  
        mock_gradle.return_value = 'gradlew'
        beam_benchmark_helper._PrebuildBeam()
  
        expected_cmd = [
            'gradlew',
            'clean',
            'assemble',
            '--stacktrace',
            '--info',
            '-p',
            'sdks/java',
            '-DintegrationTestRunner=dataflow',
            '-Dfilesystem=hdfs',
            '-Dextra_key=extra_value'
        ]
        mock_run.assert_called_once_with(expected_cmd,
                                         cwd=beam_benchmark_helper._GetBeamDir(),
                                         timeout=1500)
  
    def test_build_python_command(self):
      FLAGS.beam_python_attr = 'IT'
      FLAGS.beam_runner = 'TestRunner'
      FLAGS.beam_python_sdk_location = 'py/location.tar'
      FLAGS.beam_sdk = beam_benchmark_helper.BEAM_PYTHON_SDK
  
      with mock.patch.object(beam_benchmark_helper, '_GetGradleCommand') as mock_gradle, \
          mock.patch.object(beam_benchmark_helper, '_GetBeamDir'), \
          mock.patch.object(vm_util, 'ExecutableOnPath', return_value=True) as exec_check:
  
        mock_gradle.return_value = 'gradlew'
        mock_spec = mock.MagicMock()
        mock_spec.service_type = dpb_service.DATAFLOW
  
        actual_cmd, _ = beam_benchmark_helper.BuildBeamCommand(mock_spec,
                                                               'apache_beam.py',
                                                               ['--args'])
        expected_cmd = [
            'gradlew',
            'beam-sdks-python:integrationTest',
            '-Ptests=apache_beam.py',
            '-Pattr=IT',
            '-PpipelineOptions=--args --runner=TestRunner --sdk_location=py/location.tar',
            '--info',
            '--scan',
        ]
        self.assertListEqual(expected_cmd, actual_cmd)
        exec_check.assert_called_once()
  
    def test_build_java_gradle_command(self):
      FLAGS.beam_it_module = 'sdks/java'
      FLAGS.beam_runner = 'dataflow'
      FLAGS.beam_filesystem = 'hdfs'
      FLAGS.beam_extra_properties = '["extra_key=extra_value"]'
      FLAGS.beam_sdk = beam_benchmark_helper.BEAM_JAVA_SDK
  
      with mock.patch.object(beam_benchmark_helper, '_GetGradleCommand') as mock_gradle, \
          mock.patch.object(beam_benchmark_helper, '_GetBeamDir'), \
          mock.patch.object(vm_util, 'ExecutableOnPath', return_value=True) as exec_check:
  
        mock_gradle.return_value = 'gradlew'
        mock_spec = mock.MagicMock()
        mock_spec.service_type = dpb_service.DATAFLOW
  
        actual_cmd, _ = beam_benchmark_helper.BuildBeamCommand(mock_spec,
                                                               'org.apache.beam.sdk.java',
                                                               ['--args'])
        expected_cmd = [
            'gradlew',
            'integrationTest',
            '--tests=org.apache.beam.sdk.java',
            '-p',
            'sdks/java',
            '-DintegrationTestRunner=dataflow',
            '-Dfilesystem=hdfs',
            '-Dextra_key=extra_value',
            '-DintegrationTestPipelineOptions=[--args,"--runner=TestDataflowRunner"]',
            '--stacktrace',
            '--info',
            '--scan',
        ]
        self.assertListEqual(expected_cmd, actual_cmd)
        exec_check.assert_called_once()
  
  
  if __name__ == '__main__':
    unittest.main()
  