github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/utils/profiler_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  import os
    19  import tempfile
    20  import unittest
    21  
    22  from parameterized import param
    23  from parameterized import parameterized
    24  
    25  from apache_beam.utils.profiler import Profile
    26  
    27  
    28  class ProfilerTest(unittest.TestCase):
    29    @parameterized.expand([
    30        param(enable_cpu_memory=(True, True)),
    31        param(enable_cpu_memory=(True, False)),
    32        param(enable_cpu_memory=(False, True)),
    33        param(enable_cpu_memory=(False, False)),
    34    ])
    35    def test_profiler(self, enable_cpu_memory):
    36      try:
    37        from guppy import hpy  # pylint: disable=unused-import
    38        guppy_imported = True
    39      except ImportError:
    40        guppy_imported = False
    41  
    42      enable_cpu, enable_memory = enable_cpu_memory
    43  
    44      with tempfile.TemporaryDirectory() as tmp:
    45        with Profile('id',
    46                     profile_location=tmp,
    47                     enable_cpu_profiling=enable_cpu,
    48                     enable_memory_profiling=enable_memory,
    49                     log_results=True):
    50          hash('test')
    51        files = os.listdir(tmp)
    52        if enable_cpu:
    53          self.assertIn('cpu_profile', files)
    54        else:
    55          self.assertNotIn('cpu_profile', files)
    56  
    57        if enable_memory and guppy_imported:
    58          self.assertIn('memory_profile', files)
    59        else:
    60          self.assertNotIn('memory_profile', files)
    61  
    62  
    63  if __name__ == '__main__':
    64    unittest.main()