github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/runners/portability/job_server_test.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  # pytype: skip-file
    18  
    19  import logging
    20  import unittest
    21  
    22  from apache_beam.options.pipeline_options import PipelineOptions
    23  from apache_beam.runners.portability.job_server import JavaJarJobServer
    24  
    25  
    26  class JavaJarJobServerStub(JavaJarJobServer):
    27    def java_arguments(
    28        self, job_port, artifact_port, expansion_port, artifacts_dir):
    29      return [
    30          '--artifacts-dir',
    31          artifacts_dir,
    32          '--job-port',
    33          job_port,
    34          '--artifact-port',
    35          artifact_port,
    36          '--expansion-port',
    37          expansion_port
    38      ]
    39  
    40    def path_to_jar(self):
    41      return '/path/to/jar'
    42  
    43    @staticmethod
    44    def local_jar(url):
    45      return url
    46  
    47  
    48  class JavaJarJobServerTest(unittest.TestCase):
    49    def test_subprocess_cmd_and_endpoint(self):
    50      pipeline_options = PipelineOptions([
    51          '--job_port=8099',
    52          '--artifact_port=8098',
    53          '--expansion_port=8097',
    54          '--artifacts_dir=/path/to/artifacts/',
    55          '--job_server_java_launcher=/path/to/java',
    56          '--job_server_jvm_properties=-Dsome.property=value'
    57      ])
    58      job_server = JavaJarJobServerStub(pipeline_options)
    59      subprocess_cmd, endpoint = job_server.subprocess_cmd_and_endpoint()
    60      self.assertEqual(
    61          subprocess_cmd,
    62          [
    63              '/path/to/java',
    64              '-jar',
    65              '-Dsome.property=value',
    66              '/path/to/jar',
    67              '--artifacts-dir',
    68              '/path/to/artifacts/',
    69              '--job-port',
    70              8099,
    71              '--artifact-port',
    72              8098,
    73              '--expansion-port',
    74              8097
    75          ])
    76      self.assertEqual(endpoint, 'localhost:8099')
    77  
    78  
    79  if __name__ == '__main__':
    80    logging.getLogger().setLevel(logging.INFO)
    81    unittest.main()