1# Copyright 2023 The Pigweed Authors 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); you may not 4# use this file except in compliance with the License. You may obtain a copy of 5# the License at 6# 7# https://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, WITHOUT 11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the 12# License for the specific language governing permissions and limitations under 13# the License. 14"""Tests for pw_cli.process.""" 15 16import unittest 17import sys 18import textwrap 19 20from pw_cli import process 21 22import psutil # type: ignore 23 24# This timeout must be long enough to wait for the subprocess output, but 25# fast enough that the test doesn't take terribly long in the success case. 26FAST_TIMEOUT_SECONDS = 0.5 27KILL_SIGNALS = set({-9, 137}) 28PYTHON = sys.executable 29 30 31class RunTest(unittest.TestCase): 32 """Tests for process.run.""" 33 34 def test_returns_output(self) -> None: 35 echo_str = 'foobar' 36 print_str = f'print("{echo_str}")' 37 result = process.run(PYTHON, '-c', print_str) 38 self.assertEqual(result.output, b'foobar\n') 39 40 def test_timeout_kills_process(self) -> None: 41 sleep_100_seconds = 'import time; time.sleep(100)' 42 result = process.run( 43 PYTHON, '-c', sleep_100_seconds, timeout=FAST_TIMEOUT_SECONDS 44 ) 45 self.assertIn(result.returncode, KILL_SIGNALS) 46 47 def test_timeout_kills_subprocess(self) -> None: 48 # Spawn a subprocess which prints its pid and then waits for 100 49 # seconds. 50 sleep_in_subprocess = textwrap.dedent( 51 f""" 52 import subprocess 53 import time 54 55 child = subprocess.Popen( 56 ['{PYTHON}', '-c', 'import time; print("booh"); time.sleep(100)'] 57 ) 58 print(child.pid, flush=True) 59 time.sleep(100) 60 """ 61 ) 62 result = process.run( 63 PYTHON, '-c', sleep_in_subprocess, timeout=FAST_TIMEOUT_SECONDS 64 ) 65 self.assertIn(result.returncode, KILL_SIGNALS) 66 # The first line of the output is the PID of the child sleep process. 67 child_pid_str, sep, remainder = result.output.partition(b'\n') 68 del sep, remainder 69 child_pid = int(child_pid_str) 70 self.assertFalse(psutil.pid_exists(child_pid)) 71 72 73if __name__ == '__main__': 74 unittest.main() 75