github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/testing/test_utils_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  """Unittest for testing utilities,"""
    19  
    20  # pytype: skip-file
    21  
    22  import logging
    23  import os
    24  import tempfile
    25  import unittest
    26  
    27  import mock
    28  
    29  from apache_beam.io.filesystem import BeamIOError
    30  from apache_beam.io.filesystems import FileSystems
    31  from apache_beam.testing import test_utils as utils
    32  
    33  
    34  class TestUtilsTest(unittest.TestCase):
    35    def setUp(self):
    36      utils.patch_retry(self, utils)
    37      self.tmpdir = tempfile.mkdtemp()
    38  
    39    def test_delete_files_succeeds(self):
    40      path = os.path.join(self.tmpdir, 'f1')
    41  
    42      with open(path, 'a') as f:
    43        f.write('test')
    44  
    45      assert FileSystems.exists(path)
    46      utils.delete_files([path])
    47      assert not FileSystems.exists(path)
    48  
    49    def test_delete_files_fails_with_io_error(self):
    50      path = os.path.join(self.tmpdir, 'f2')
    51  
    52      with self.assertRaises(BeamIOError) as error:
    53        utils.delete_files([path])
    54      self.assertTrue(
    55          error.exception.args[0].startswith('Delete operation failed'))
    56      self.assertEqual(list(error.exception.exception_details.keys()), [path])
    57  
    58    def test_delete_files_fails_with_invalid_arg(self):
    59      with self.assertRaises(RuntimeError):
    60        utils.delete_files([])
    61  
    62    def test_temp_dir_removes_files(self):
    63      with utils.TempDir() as tempdir:
    64        dir_path = tempdir.get_path()
    65        file_path = tempdir.create_temp_file()
    66        self.assertTrue(os.path.exists(dir_path))
    67        self.assertTrue(os.path.exists(file_path))
    68  
    69      self.assertFalse(os.path.exists(dir_path))
    70      self.assertFalse(os.path.exists(file_path))
    71  
    72    def test_temp_file_field_correct(self):
    73      with utils.TempDir() as tempdir:
    74        filename = tempdir.create_temp_file(
    75            suffix='.txt', lines=[b'line1\n', b'line2\n', b'line3\n'])
    76        self.assertTrue(filename.endswith('.txt'))
    77  
    78        with open(filename, 'rb') as f:
    79          self.assertEqual(f.readline(), b'line1\n')
    80          self.assertEqual(f.readline(), b'line2\n')
    81          self.assertEqual(f.readline(), b'line3\n')
    82  
    83    def test_cleanup_subscriptions(self):
    84      sub_client = mock.Mock()
    85      sub = mock.Mock()
    86      sub.name = 'test_sub'
    87      utils.cleanup_subscriptions(sub_client, [sub])
    88      sub_client.delete_subscription.assert_called_with(subscription=sub.name)
    89  
    90    def test_cleanup_topics(self):
    91      pub_client = mock.Mock()
    92      topic = mock.Mock()
    93      topic.name = 'test_topic'
    94      utils.cleanup_topics(pub_client, [topic])
    95      pub_client.delete_topic.assert_called_with(topic=topic.name)
    96  
    97  
    98  if __name__ == '__main__':
    99    logging.getLogger().setLevel(logging.INFO)
   100    unittest.main()