def test_build_python_gradle_command(self): FLAGS.beam_python_attr = 'IT' FLAGS.beam_it_module = ':sdks:python' 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', ':sdks:python:integrationTest', '-Dtests=apache_beam.py', '-Dattr=IT', '-DpipelineOptions=--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', ':sdks:java:integrationTest', '--tests=org.apache.beam.sdk.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()
def SubmitJob(self, jarfile, classname, job_poll_interval=None, job_arguments=None, job_stdout_file=None, job_type=None): """See base class.""" if job_type == self.BEAM_JOB_TYPE: full_cmd, base_dir = beam_benchmark_helper.BuildBeamCommand( self.spec, classname, job_arguments) stdout, _, retcode = vm_util.IssueCommand( full_cmd, cwd=base_dir, timeout=FLAGS.beam_it_timeout, raise_on_failure=False) assert retcode == 0, "Integration Test Failed." return worker_machine_type = self.spec.worker_group.vm_spec.machine_type num_workers = self.spec.worker_count max_num_workers = self.spec.worker_count if self.spec.worker_group.disk_spec and \ self.spec.worker_group.disk_spec.disk_size: disk_size_gb = self.spec.worker_group.disk_spec.disk_size elif self.spec.worker_group.vm_spec.boot_disk_size: disk_size_gb = self.spec.worker_group.vm_spec.boot_disk_size else: disk_size_gb = None cmd = [] # Needed to verify java executable is on the path dataflow_executable = 'java' if not vm_util.ExecutableOnPath(dataflow_executable): raise errors.Setup.MissingExecutableError( 'Could not find required executable "%s"' % dataflow_executable) cmd.append(dataflow_executable) cmd.append('-cp') cmd.append(jarfile) cmd.append(classname) cmd += job_arguments cmd.append('--workerMachineType={}'.format(worker_machine_type)) cmd.append('--numWorkers={}'.format(num_workers)) cmd.append('--maxNumWorkers={}'.format(max_num_workers)) if disk_size_gb: cmd.append('--diskSizeGb={}'.format(disk_size_gb)) cmd.append('--defaultWorkerLogLevel={}'.format(FLAGS.dpb_log_level)) stdout, _, _ = vm_util.IssueCommand(cmd)