1# Copyright 2024, The Android Open Source Project 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15"""Testing utilities for tests in the CI package.""" 16 17import logging 18import os 19import unittest 20import subprocess 21import pathlib 22import shutil 23import tempfile 24 25 26# Export the TestCase class to reduce the number of imports tests have to list. 27TestCase = unittest.TestCase 28 29 30def process_alive(pid): 31 """Check For the existence of a pid.""" 32 33 try: 34 os.kill(pid, 0) 35 except OSError: 36 return False 37 38 return True 39 40 41class TemporaryProcessSession: 42 43 def __init__(self, test_case: TestCase): 44 self._created_processes = [] 45 test_case.addCleanup(self.cleanup) 46 47 def create(self, args, kwargs): 48 p = subprocess.Popen(*args, **kwargs, start_new_session=True) 49 self._created_processes.append(p) 50 return p 51 52 def cleanup(self): 53 for p in self._created_processes: 54 if not process_alive(p.pid): 55 return 56 os.killpg(os.getpgid(p.pid), signal.SIGKILL) 57 58 59class TestTemporaryDirectory: 60 61 def __init__(self, delete: bool, ): 62 self._delete = delete 63 64 @classmethod 65 def create(cls, test_case: TestCase, delete: bool = True): 66 temp_dir = TestTemporaryDirectory(delete) 67 temp_dir._dir = pathlib.Path(tempfile.mkdtemp()) 68 test_case.addCleanup(temp_dir.cleanup) 69 return temp_dir._dir 70 71 def get_dir(self): 72 return self._dir 73 74 def cleanup(self): 75 if not self._delete: 76 return 77 shutil.rmtree(self._dir, ignore_errors=True) 78 79 80def main(): 81 82 # Disable logging since it breaks the TF Python test output parser. 83 # TODO(hzalek): Use TF's `test-output-file` option to re-enable logging. 84 logging.getLogger().disabled = True 85 86 unittest.main() 87