scripts/lib/async_utils/tests/async_utils_test.py - fuchsia - Git at Google

 # Copyright 2023 The Fuchsia Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.

 import asyncio
 import multiprocessing
 import os
 import signal
 import stat
 import tempfile
 import unittest

 import async_utils.signals
 from async_utils import command


 class TestCommand(unittest.IsolatedAsyncioTestCase):
     def assertStdout(self, event: command.CommandEvent, line: bytes) -> None:
         """Helper to assert on contents of a StdoutEvent.

         Args:
             event (command.CommandEvent): Event to cast and compare.
             line (bytes): Expected line value.
         """
         self.assertTrue(isinstance(event, command.StdoutEvent))
         assert isinstance(event, command.StdoutEvent)
         e: command.StdoutEvent = event
         self.assertEqual(e.text, line)

     def assertStderr(self, event: command.CommandEvent, line: bytes) -> None:
         """Helper to assert on contents of a StderrEvent.

         Args:
             event (command.CommandEvent): Event to cast and compare.
             line (bytes): Expected line value.
         """
         self.assertTrue(isinstance(event, command.StderrEvent))
         assert isinstance(event, command.StderrEvent)
         e: command.StderrEvent = event
         self.assertEqual(e.text, line)

     def assertTermination(
         self, event: command.CommandEvent, return_code: int
     ) -> None:
         """Helper to assert on contents of a TerminationEvent.

         Args:
             event (command.CommandEvent): Event to cast and compare.
             return_code (int): Expected return code.
         """
         self.assertTrue(isinstance(event, command.TerminationEvent))
         assert isinstance(event, command.TerminationEvent)
         e: command.TerminationEvent = event
         self.assertEqual(e.return_code, return_code)

     async def test_basic_command(self) -> None:
         """Test running a basic command and getting the output.

         We create a file in a temporary directory and simply assert that `ls`
         prints that file as output.
         """
         with tempfile.TemporaryDirectory() as td:
             with open(os.path.join(td, "temp-file.txt"), "w") as f:
                 f.write("hello world")

             cmd = await command.AsyncCommand.create("ls", ".", env={"CWD": td})
             events = []
             complete = await cmd.run_to_completion(
                 lambda event: events.append(event)
             )
             self.assertEqual(len(events), 2, f"Events was actually {events}")

             self.assertStdout(events[0], b"temp-file.txt\n")
             self.assertTermination(events[1], 0)

             self.assertEqual(complete.stdout, "temp-file.txt\n")
             self.assertEqual(complete.return_code, 0)

     async def test_command_with_input(self) -> None:
         """Test passing input to a command."""
         cmd = await command.AsyncCommand.create(
             "cat", input_bytes=b"hello\nworld"
         )
         events = []
         complete = await cmd.run_to_completion(
             lambda event: events.append(event)
         )
         self.assertEqual(len(events), 3, f"Events was actually {events}")

         self.assertStdout(events[0], b"hello\n")
         self.assertStdout(events[1], b"world")
         self.assertTermination(events[2], 0)

         self.assertEqual(complete.stdout, "hello\nworld")
         self.assertEqual(complete.return_code, 0)

     async def test_wrapper_with_input(self) -> None:
         """Test passing input to a command that uses a symbolizer"""
         cmd = await command.AsyncCommand.create(
             "cat",
             symbolizer_args=["sed", "s/hello/goodbye/g"],
             input_bytes=b"hello\nworld",
         )
         events = []
         complete = await cmd.run_to_completion(
             lambda event: events.append(event)
         )
         self.assertEqual(len(events), 3, f"Events was actually {events}")

         self.assertStdout(events[0], b"goodbye\n")
         self.assertStdout(events[1], b"world")
         self.assertTermination(events[2], 0)

         self.assertEqual(complete.stdout, "goodbye\nworld")
         self.assertEqual(complete.return_code, 0)

     async def test_long_line_output(self) -> None:
         """Test processing a very large output from a program"""
         cmd = await command.AsyncCommand.create(
             "cat",
             input_bytes=b"a" * 1024 * 1024,
         )
         events = []
         complete = await cmd.run_to_completion(
             lambda event: events.append(event)
         )
         self.assertEqual(len(events), 2, f"Events was actually {events}")

         self.assertStdout(events[0], b"a" * 1024 * 1024)
         self.assertTermination(events[1], 0)

         self.assertEqual(complete.stdout, "a" * 1024 * 1024)
         self.assertEqual(complete.return_code, 0)

     async def test_basic_command_with_long_timeout(self) -> None:
         """Test running a basic command and getting the output.

         We create a file in a temporary directory and simply assert that `ls`
         prints that file as output.
         """
         with tempfile.TemporaryDirectory() as td:
             with open(os.path.join(td, "temp-file.txt"), "w") as f:
                 f.write("hello world")

             cmd = await command.AsyncCommand.create(
                 "ls", ".", env={"CWD": td}, timeout=3600
             )
             events = []
             complete = await cmd.run_to_completion(
                 lambda event: events.append(event)
             )
             self.assertEqual(len(events), 2, f"Events was actually {events}")

             self.assertStdout(events[0], b"temp-file.txt\n")
             self.assertTermination(events[1], 0)

             self.assertEqual(complete.stdout, "temp-file.txt\n")
             self.assertEqual(complete.return_code, 0)
             self.assertFalse(complete.was_timeout)

     async def test_with_stderr(self) -> None:
         """Test running a command with stderr output.

         We create a temporary directory and try to `ls` a file we know does not
         exist. `ls` should print to stderr and report an error return code.
         """
         with tempfile.TemporaryDirectory() as td:
             cmd = await command.AsyncCommand.create(
                 "ls", os.path.join(td, "does-not-exist")
             )
             complete = await cmd.run_to_completion()
             self.assertEqual(complete.stdout, "")
             self.assertNotEqual(complete.stderr, "")
             self.assertNotEqual(complete.return_code, 0)

     async def test_symbolized_command(self) -> None:
         """Test piping output through another program.

         We run `ls` as in the above tests, but this time we pipe the output
         through `sed` to change the word "temp" to "temporary" and assert on
         the new output.
         """
         with tempfile.TemporaryDirectory() as td:
             with open(os.path.join(td, "temp-file.txt"), "w") as f:
                 f.write("hello world")

             cmd = await command.AsyncCommand.create(
                 "ls",
                 ".",
                 env={"CWD": td},
                 symbolizer_args=["sed", "s/temp/temporary/g"],
             )
             events = []
             await cmd.run_to_completion(lambda event: events.append(event))
             self.assertEqual(len(events), 2, f"Events was actually {events}")

             self.assertStdout(events[0], b"temporary-file.txt\n")
             self.assertTermination(events[1], 0)

     async def test_terminate_and_kill(self) -> None:
         """Test that we can terminate and kill programs.

         We spawn `sleep` to run for over a day, then terminate it. We expect
         the return code to be set by the OS to represent that the program
         was killed.
         """
         cmd = await command.AsyncCommand.create("sleep", "100000")
         task = asyncio.create_task(cmd.run_to_completion())
         cmd.terminate()
         out: command.CommandOutput = await task
         self.assertEqual(out.return_code, -15)

         cmd = await command.AsyncCommand.create("sleep", "100000")
         task = asyncio.create_task(cmd.run_to_completion())
         cmd.kill()
         out = await task
         self.assertEqual(out.return_code, -9)

         cmd = await command.AsyncCommand.create(
             "sleep", "100000", symbolizer_args=["sleep", "100000"]
         )
         task = asyncio.create_task(cmd.run_to_completion())
         cmd.terminate()
         out = await task
         self.assertEqual(out.return_code, -15)
         self.assertEqual(out.wrapper_return_code, -15)

         cmd = await command.AsyncCommand.create(
             "sleep", "100000", symbolizer_args=["sleep", "100000"]
         )
         task = asyncio.create_task(cmd.run_to_completion())
         cmd.kill()
         out = await task
         self.assertEqual(out.return_code, -9)
         self.assertEqual(out.wrapper_return_code, -9)

     async def test_kill_process_groups(self) -> None:
         """Test that terminating a program kills the entire process group."""

         BASH_SHORT = "#!/usr/bin/env bash\nsleep .1\necho 'OK'"
         BASH_LONG = "#!/usr/bin/env bash\nsleep 100000\necho 'OK'"
         with tempfile.TemporaryDirectory() as td:
             # Create scripts and make them executable.
             paths = [
                 os.path.join(td, "short.sh"),
                 os.path.join(td, "long.sh"),
             ]
             short_path, long_path = paths
             with open(short_path, "w") as f:
                 f.write(BASH_SHORT)
             with open(long_path, "w") as f:
                 f.write(BASH_LONG)
             for name in paths:
                 st = os.stat(name)
                 os.chmod(
                     name,
                     st.st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH,
                 )

             # Make sure we can run the commands to start with, and they do not hang.
             cmd = await command.AsyncCommand.create(
                 short_path,
                 env={"CWD": td},
                 symbolizer_args=[short_path],
             )
             events = []
             await cmd.run_to_completion(lambda event: events.append(event))
             self.assertEqual(len(events), 2, f"Events was actually {events}")

             self.assertStdout(events[0], b"OK\n")
             # This causes spurious failures on Mac for some reason, where
             # the return value is sometimes -13.
             # self.assertTermination(events[1], 0)

             # Run the long-running shell script, and ensure terminating it does not hang.
             cmd = await command.AsyncCommand.create(
                 long_path,
                 env={"CWD": td},
                 symbolizer_args=[long_path],
             )
             await asyncio.sleep(0.001)
             cmd.terminate()
             events = []
             await cmd.run_to_completion(lambda event: events.append(event))
             self.assertEqual(len(events), 1, f"Events was actually {events}")
             self.assertTermination(events[0], -15)

             # Run again, this time using SIGKILL.
             cmd = await command.AsyncCommand.create(
                 long_path,
                 env={"CWD": td},
                 symbolizer_args=[long_path],
             )
             await asyncio.sleep(0.001)
             cmd.kill()
             events = []
             await cmd.run_to_completion(lambda event: events.append(event))
             self.assertEqual(len(events), 1, f"Events was actually {events}")
             self.assertTermination(events[0], -9)

     async def test_timeout(self) -> None:
         """Test that commands timeout"""
         cmd = await command.AsyncCommand.create("sleep", "100000", timeout=0.1)
         task = asyncio.create_task(cmd.run_to_completion())
         out: command.CommandOutput = await task
         self.assertEqual(out.return_code, -15)
         self.assertTrue(out.was_timeout)

     def test_invalid_program(self) -> None:
         """Test running a program that doesn't exist, and expect an error."""
         self.assertRaises(
             command.AsyncCommandError,
             lambda: asyncio.run(command.AsyncCommand.create("..........")),
         )


 class TestSignals(unittest.TestCase):
     def test_async_signal_handler(self) -> None:
         """Test that registered signal handlers work appropriately."""

         multiprocessing.set_start_method("fork", force=True)

         output_directory = tempfile.TemporaryDirectory()
         self.addCleanup(output_directory.cleanup)
         output_file_name = os.path.join(output_directory.name, "output.txt")

         def main(output_file_name: str) -> None:
             async def internal_main() -> None:
                 os.kill(os.getpid(), signal.SIGTERM)
                 await asyncio.sleep(120)

             asyncio.set_event_loop(asyncio.new_event_loop())
             loop = asyncio.get_event_loop()

             fut = asyncio.ensure_future(internal_main())

             def write_output() -> None:
                 with open(output_file_name, "a") as f:
                     f.write("Handler printed message\n")
                 fut.cancel()

             async_utils.signals.register_on_terminate_signal(write_output)
             try:
                 loop.run_until_complete(fut)
             except asyncio.CancelledError:
                 with open(output_file_name, "a") as f:
                     f.write("Cancelled\n")

         proc = multiprocessing.Process(target=main, args=(output_file_name,))
         proc.start()
         proc.join()

         lines: list[str]
         with open(output_file_name, "r") as f:
             lines = [line.strip() for line in f.readlines()]

         self.assertListEqual(lines, ["Handler printed message", "Cancelled"])


 if __name__ == "__main__":
     unittest.main()
	# Copyright 2023 The Fuchsia Authors. All rights reserved.
	# Use of this source code is governed by a BSD-style license that can be
	# found in the LICENSE file.

	import asyncio
	import multiprocessing
	import os
	import signal
	import stat
	import tempfile
	import unittest

	import async_utils.signals
	from async_utils import command


	class TestCommand(unittest.IsolatedAsyncioTestCase):
	def assertStdout(self, event: command.CommandEvent, line: bytes) -> None:
	"""Helper to assert on contents of a StdoutEvent.

	Args:
	event (command.CommandEvent): Event to cast and compare.
	line (bytes): Expected line value.
	"""
	self.assertTrue(isinstance(event, command.StdoutEvent))
	assert isinstance(event, command.StdoutEvent)
	e: command.StdoutEvent = event
	self.assertEqual(e.text, line)

	def assertStderr(self, event: command.CommandEvent, line: bytes) -> None:
	"""Helper to assert on contents of a StderrEvent.

	Args:
	event (command.CommandEvent): Event to cast and compare.
	line (bytes): Expected line value.
	"""
	self.assertTrue(isinstance(event, command.StderrEvent))
	assert isinstance(event, command.StderrEvent)
	e: command.StderrEvent = event
	self.assertEqual(e.text, line)

	def assertTermination(
	self, event: command.CommandEvent, return_code: int
	) -> None:
	"""Helper to assert on contents of a TerminationEvent.

	Args:
	event (command.CommandEvent): Event to cast and compare.
	return_code (int): Expected return code.
	"""
	self.assertTrue(isinstance(event, command.TerminationEvent))
	assert isinstance(event, command.TerminationEvent)
	e: command.TerminationEvent = event
	self.assertEqual(e.return_code, return_code)

	async def test_basic_command(self) -> None:
	"""Test running a basic command and getting the output.

	We create a file in a temporary directory and simply assert that `ls`
	prints that file as output.
	"""
	with tempfile.TemporaryDirectory() as td:
	with open(os.path.join(td, "temp-file.txt"), "w") as f:
	f.write("hello world")

	cmd = await command.AsyncCommand.create("ls", ".", env={"CWD": td})
	events = []
	complete = await cmd.run_to_completion(
	lambda event: events.append(event)
	)
	self.assertEqual(len(events), 2, f"Events was actually {events}")

	self.assertStdout(events[0], b"temp-file.txt\n")
	self.assertTermination(events[1], 0)

	self.assertEqual(complete.stdout, "temp-file.txt\n")
	self.assertEqual(complete.return_code, 0)

	async def test_command_with_input(self) -> None:
	"""Test passing input to a command."""
	cmd = await command.AsyncCommand.create(
	"cat", input_bytes=b"hello\nworld"
	)
	events = []
	complete = await cmd.run_to_completion(
	lambda event: events.append(event)
	)
	self.assertEqual(len(events), 3, f"Events was actually {events}")

	self.assertStdout(events[0], b"hello\n")
	self.assertStdout(events[1], b"world")
	self.assertTermination(events[2], 0)

	self.assertEqual(complete.stdout, "hello\nworld")
	self.assertEqual(complete.return_code, 0)

	async def test_wrapper_with_input(self) -> None:
	"""Test passing input to a command that uses a symbolizer"""
	cmd = await command.AsyncCommand.create(
	"cat",
	symbolizer_args=["sed", "s/hello/goodbye/g"],
	input_bytes=b"hello\nworld",
	)
	events = []
	complete = await cmd.run_to_completion(
	lambda event: events.append(event)
	)
	self.assertEqual(len(events), 3, f"Events was actually {events}")

	self.assertStdout(events[0], b"goodbye\n")
	self.assertStdout(events[1], b"world")
	self.assertTermination(events[2], 0)

	self.assertEqual(complete.stdout, "goodbye\nworld")
	self.assertEqual(complete.return_code, 0)

	async def test_long_line_output(self) -> None:
	"""Test processing a very large output from a program"""
	cmd = await command.AsyncCommand.create(
	"cat",
	input_bytes=b"a" * 1024 * 1024,
	)
	events = []
	complete = await cmd.run_to_completion(
	lambda event: events.append(event)
	)
	self.assertEqual(len(events), 2, f"Events was actually {events}")

	self.assertStdout(events[0], b"a" * 1024 * 1024)
	self.assertTermination(events[1], 0)

	self.assertEqual(complete.stdout, "a" * 1024 * 1024)
	self.assertEqual(complete.return_code, 0)

	async def test_basic_command_with_long_timeout(self) -> None:
	"""Test running a basic command and getting the output.

	We create a file in a temporary directory and simply assert that `ls`
	prints that file as output.
	"""
	with tempfile.TemporaryDirectory() as td:
	with open(os.path.join(td, "temp-file.txt"), "w") as f:
	f.write("hello world")

	cmd = await command.AsyncCommand.create(
	"ls", ".", env={"CWD": td}, timeout=3600
	)
	events = []
	complete = await cmd.run_to_completion(
	lambda event: events.append(event)
	)
	self.assertEqual(len(events), 2, f"Events was actually {events}")

	self.assertStdout(events[0], b"temp-file.txt\n")
	self.assertTermination(events[1], 0)

	self.assertEqual(complete.stdout, "temp-file.txt\n")
	self.assertEqual(complete.return_code, 0)
	self.assertFalse(complete.was_timeout)

	async def test_with_stderr(self) -> None:
	"""Test running a command with stderr output.

	We create a temporary directory and try to `ls` a file we know does not
	exist. `ls` should print to stderr and report an error return code.
	"""
	with tempfile.TemporaryDirectory() as td:
	cmd = await command.AsyncCommand.create(
	"ls", os.path.join(td, "does-not-exist")
	)
	complete = await cmd.run_to_completion()
	self.assertEqual(complete.stdout, "")
	self.assertNotEqual(complete.stderr, "")
	self.assertNotEqual(complete.return_code, 0)

	async def test_symbolized_command(self) -> None:
	"""Test piping output through another program.

	We run `ls` as in the above tests, but this time we pipe the output
	through `sed` to change the word "temp" to "temporary" and assert on
	the new output.
	"""
	with tempfile.TemporaryDirectory() as td:
	with open(os.path.join(td, "temp-file.txt"), "w") as f:
	f.write("hello world")

	cmd = await command.AsyncCommand.create(
	"ls",
	".",
	env={"CWD": td},
	symbolizer_args=["sed", "s/temp/temporary/g"],
	)
	events = []
	await cmd.run_to_completion(lambda event: events.append(event))
	self.assertEqual(len(events), 2, f"Events was actually {events}")

	self.assertStdout(events[0], b"temporary-file.txt\n")
	self.assertTermination(events[1], 0)

	async def test_terminate_and_kill(self) -> None:
	"""Test that we can terminate and kill programs.

	We spawn `sleep` to run for over a day, then terminate it. We expect
	the return code to be set by the OS to represent that the program
	was killed.
	"""
	cmd = await command.AsyncCommand.create("sleep", "100000")
	task = asyncio.create_task(cmd.run_to_completion())
	cmd.terminate()
	out: command.CommandOutput = await task
	self.assertEqual(out.return_code, -15)

	cmd = await command.AsyncCommand.create("sleep", "100000")
	task = asyncio.create_task(cmd.run_to_completion())
	cmd.kill()
	out = await task
	self.assertEqual(out.return_code, -9)

	cmd = await command.AsyncCommand.create(
	"sleep", "100000", symbolizer_args=["sleep", "100000"]
	)
	task = asyncio.create_task(cmd.run_to_completion())
	cmd.terminate()
	out = await task
	self.assertEqual(out.return_code, -15)
	self.assertEqual(out.wrapper_return_code, -15)

	cmd = await command.AsyncCommand.create(
	"sleep", "100000", symbolizer_args=["sleep", "100000"]
	)
	task = asyncio.create_task(cmd.run_to_completion())
	cmd.kill()
	out = await task
	self.assertEqual(out.return_code, -9)
	self.assertEqual(out.wrapper_return_code, -9)

	async def test_kill_process_groups(self) -> None:
	"""Test that terminating a program kills the entire process group."""

	BASH_SHORT = "#!/usr/bin/env bash\nsleep .1\necho 'OK'"
	BASH_LONG = "#!/usr/bin/env bash\nsleep 100000\necho 'OK'"
	with tempfile.TemporaryDirectory() as td:
	# Create scripts and make them executable.
	paths = [
	os.path.join(td, "short.sh"),
	os.path.join(td, "long.sh"),
	]
	short_path, long_path = paths
	with open(short_path, "w") as f:
	f.write(BASH_SHORT)
	with open(long_path, "w") as f:
	f.write(BASH_LONG)
	for name in paths:
	st = os.stat(name)
	os.chmod(
	name,
	st.st_mode \| stat.S_IXUSR \| stat.S_IXGRP \| stat.S_IXOTH,
	)

	# Make sure we can run the commands to start with, and they do not hang.
	cmd = await command.AsyncCommand.create(
	short_path,
	env={"CWD": td},
	symbolizer_args=[short_path],
	)
	events = []
	await cmd.run_to_completion(lambda event: events.append(event))
	self.assertEqual(len(events), 2, f"Events was actually {events}")

	self.assertStdout(events[0], b"OK\n")
	# This causes spurious failures on Mac for some reason, where
	# the return value is sometimes -13.
	# self.assertTermination(events[1], 0)

	# Run the long-running shell script, and ensure terminating it does not hang.
	cmd = await command.AsyncCommand.create(
	long_path,
	env={"CWD": td},
	symbolizer_args=[long_path],
	)
	await asyncio.sleep(0.001)
	cmd.terminate()
	events = []
	await cmd.run_to_completion(lambda event: events.append(event))
	self.assertEqual(len(events), 1, f"Events was actually {events}")
	self.assertTermination(events[0], -15)

	# Run again, this time using SIGKILL.
	cmd = await command.AsyncCommand.create(
	long_path,
	env={"CWD": td},
	symbolizer_args=[long_path],
	)
	await asyncio.sleep(0.001)
	cmd.kill()
	events = []
	await cmd.run_to_completion(lambda event: events.append(event))
	self.assertEqual(len(events), 1, f"Events was actually {events}")
	self.assertTermination(events[0], -9)

	async def test_timeout(self) -> None:
	"""Test that commands timeout"""
	cmd = await command.AsyncCommand.create("sleep", "100000", timeout=0.1)
	task = asyncio.create_task(cmd.run_to_completion())
	out: command.CommandOutput = await task
	self.assertEqual(out.return_code, -15)
	self.assertTrue(out.was_timeout)

	def test_invalid_program(self) -> None:
	"""Test running a program that doesn't exist, and expect an error."""
	self.assertRaises(
	command.AsyncCommandError,
	lambda: asyncio.run(command.AsyncCommand.create("..........")),
	)


	class TestSignals(unittest.TestCase):
	def test_async_signal_handler(self) -> None:
	"""Test that registered signal handlers work appropriately."""

	multiprocessing.set_start_method("fork", force=True)

	output_directory = tempfile.TemporaryDirectory()
	self.addCleanup(output_directory.cleanup)
	output_file_name = os.path.join(output_directory.name, "output.txt")

	def main(output_file_name: str) -> None:
	async def internal_main() -> None:
	os.kill(os.getpid(), signal.SIGTERM)
	await asyncio.sleep(120)

	asyncio.set_event_loop(asyncio.new_event_loop())
	loop = asyncio.get_event_loop()

	fut = asyncio.ensure_future(internal_main())

	def write_output() -> None:
	with open(output_file_name, "a") as f:
	f.write("Handler printed message\n")
	fut.cancel()

	async_utils.signals.register_on_terminate_signal(write_output)
	try:
	loop.run_until_complete(fut)
	except asyncio.CancelledError:
	with open(output_file_name, "a") as f:
	f.write("Cancelled\n")

	proc = multiprocessing.Process(target=main, args=(output_file_name,))
	proc.start()
	proc.join()

	lines: list[str]
	with open(output_file_name, "r") as f:
	lines = [line.strip() for line in f.readlines()]

	self.assertListEqual(lines, ["Handler printed message", "Cancelled"])


	if __name__ == "__main__":
	unittest.main()