github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/examples/complete/distribopt_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  
    18  """Test for the distrib_optimization example."""
    19  
    20  # pytype: skip-file
    21  
    22  import logging
    23  import unittest
    24  import uuid
    25  from ast import literal_eval as make_tuple
    26  
    27  import numpy as np
    28  import pytest
    29  from mock import MagicMock
    30  from mock import patch
    31  
    32  from apache_beam.testing.test_pipeline import TestPipeline
    33  from apache_beam.testing.test_utils import create_file
    34  from apache_beam.testing.test_utils import read_files_from_pattern
    35  
    36  FILE_CONTENTS = 'OP01,8,12,0,12\n' \
    37                  'OP02,30,14,3,12\n' \
    38                  'OP03,25,7,3,14\n' \
    39                  'OP04,87,7,2,2\n' \
    40                  'OP05,19,1,7,10'
    41  
    42  EXPECTED_MAPPING = {
    43      'OP01': 'A', 'OP02': 'B', 'OP03': 'B', 'OP04': 'C', 'OP05': 'A'
    44  }
    45  
    46  
    47  class DistribOptimizationTest(unittest.TestCase):
    48    #TODO(https://github.com/apache/beam/issues/23606) Fix and enable
    49    @pytest.mark.sickbay_dataflow
    50    @pytest.mark.examples_postcommit
    51    def test_basics(self):
    52      test_pipeline = TestPipeline(is_integration_test=True)
    53  
    54      # Setup the files with expected content.
    55      temp_location = test_pipeline.get_option('temp_location')
    56      input = '/'.join([temp_location, str(uuid.uuid4()), 'input.txt'])
    57      output = '/'.join([temp_location, str(uuid.uuid4()), 'result'])
    58      create_file(input, FILE_CONTENTS)
    59      extra_opts = {'input': input, 'output': output}
    60  
    61      # Run pipeline
    62      # Avoid dependency on SciPy
    63      scipy_mock = MagicMock()
    64      result_mock = MagicMock(x=np.ones(3))
    65      scipy_mock.optimize.minimize = MagicMock(return_value=result_mock)
    66      modules = {'scipy': scipy_mock, 'scipy.optimize': scipy_mock.optimize}
    67  
    68      with patch.dict('sys.modules', modules):
    69        from apache_beam.examples.complete import distribopt
    70        distribopt.run(
    71            test_pipeline.get_full_options_as_args(**extra_opts),
    72            save_main_session=False)
    73  
    74      # Load result file and compare.
    75      lines = read_files_from_pattern('%s*' % output).splitlines()
    76  
    77      # Only 1 result
    78      self.assertEqual(len(lines), 1)
    79  
    80      # parse result line and verify optimum
    81      optimum = make_tuple(lines[0])
    82      self.assertAlmostEqual(optimum['cost'], 454.39597, places=3)
    83      self.assertDictEqual(optimum['mapping'], EXPECTED_MAPPING)
    84      production = optimum['production']
    85      for plant in ['A', 'B', 'C']:
    86        np.testing.assert_almost_equal(production[plant], np.ones(3))
    87  
    88  
    89  if __name__ == '__main__':
    90    logging.getLogger().setLevel(logging.INFO)
    91    unittest.main()