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)