xref: /aosp_15_r20/external/toolchain-utils/pgo_tools/pgo_tools_test.py (revision 760c253c1ed00ce9abd48f8546f08516e57485fe)
1#!/usr/bin/env python3
2# Copyright 2023 The ChromiumOS Authors
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""Tests for pgo_tools."""
7
8from pathlib import Path
9import textwrap
10import unittest
11from unittest import mock
12
13import pgo_tools
14
15
16class Test(unittest.TestCase):
17    """Tests for pgo_tools."""
18
19    @mock.patch.object(pgo_tools, "run")
20    def test_pgo_generate_checking_works(self, mock_run):
21        equery_u_output = textwrap.dedent(
22            """\
23            [ Legend : U - final flag setting for installation]
24            [        : I - package is installed with flag     ]
25            [ Colors : set, unset                             ]
26             * Found these USE flags for sys-devel/llvm:
27             U I
28             - - llvm-next_pgo_use              : <unknown>
29             - - llvm-tot                       : <unknown>
30             - + llvm_pgo_generate              : <unknown>
31             + - llvm_pgo_use                   : <unknown>
32             """
33        )
34        mock_run.return_value.stdout = equery_u_output
35        self.assertTrue(pgo_tools.installed_llvm_has_pgo_generate_enabled())
36
37        mock_run.assert_called_once()
38
39        mock_run.return_value.stdout = equery_u_output.replace(
40            "+ llvm_pgo_generate", "- llvm_pgo_generate"
41        )
42        self.assertFalse(pgo_tools.installed_llvm_has_pgo_generate_enabled())
43
44    @mock.patch.object(pgo_tools, "run")
45    def test_pgo_generate_checking_raises_on_zero_pgo_updates(self, mock_run):
46        mock_run.return_value.stdout = textwrap.dedent(
47            """\
48            [ Legend : U - final flag setting for installation]
49             - - llvm-next_pgo_use              : <unknown>
50             - - llvm-tot                       : <unknown>
51             + - llvm_pgo_use                   : <unknown>
52             """
53        )
54        with self.assertRaisesRegex(ValueError, "^No llvm_pgo_generate"):
55            pgo_tools.installed_llvm_has_pgo_generate_enabled()
56
57    @mock.patch.object(pgo_tools, "run")
58    def test_pgo_generate_checking_raises_on_many_pgo_updates(self, mock_run):
59        mock_run.return_value.stdout = textwrap.dedent(
60            """\
61            [ Legend : U - final flag setting for installation]
62             - - llvm-next_pgo_use              : <unknown>
63             - - llvm-tot                       : <unknown>
64             - + llvm_pgo_generate              : <unknown>
65             - + llvm_pgo_generate              : <unknown>
66             + - llvm_pgo_use                   : <unknown>
67             """
68        )
69        with self.assertRaisesRegex(ValueError, "^Multiple llvm_pgo_generate"):
70            pgo_tools.installed_llvm_has_pgo_generate_enabled()
71
72    @mock.patch.object(pgo_tools, "run")
73    def test_pgo_generate_ignores_nonexact_use_flags(self, mock_run):
74        mock_run.return_value.stdout = textwrap.dedent(
75            """\
76            [ Legend : U - final flag setting for installation]
77             - - llvm-next_pgo_use              : <unknown>
78             - - llvm-tot                       : <unknown>
79             - + llvm_pgo_generate              : <unknown>
80             - - llvm_pgo_generate2             : <unknown>
81             - - 2llvm_pgo_generate             : <unknown>
82             + - llvm_pgo_use                   : <unknown>
83             """
84        )
85        self.assertTrue(pgo_tools.installed_llvm_has_pgo_generate_enabled())
86
87    def test_quickpkg_restoration_works(self):
88        self.assertEqual(
89            pgo_tools.generate_quickpkg_restoration_command(
90                Path("/path/to/sys-devel/llvm-1234-r1.tbz2")
91            ),
92            ["sudo", "emerge", "--usepkgonly", "=sys-devel/llvm-1234-r1"],
93        )
94
95    def test_temporary_file_creation_works(self):
96        with pgo_tools.temporary_file("foo_bar_") as tmp:
97            self.assertTrue(tmp.name.startswith("foo_bar_"), tmp.name)
98            self.assertTrue(tmp.exists())
99        self.assertFalse(tmp.exists())
100
101    def test_temporary_file_deletion_is_fine_if_file_does_not_exist(self):
102        # This test ensures this `with`'s `__exit__` block doesn't `raise`.
103        with pgo_tools.temporary_file("foo_bar_") as tmp:
104            tmp.unlink()
105
106
107if __name__ == "__main__":
108    unittest.main()
109