xref: /aosp_15_r20/external/toolchain-utils/afdo_tools/update_kernel_afdo_test.py (revision 760c253c1ed00ce9abd48f8546f08516e57485fe)
1#!/usr/bin/env python3
2# Copyright 2024 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 update_kernel_afdo."""
7
8import datetime
9from pathlib import Path
10import shutil
11import subprocess
12import tempfile
13import textwrap
14import unittest
15from unittest import mock
16
17import update_kernel_afdo
18
19
20class Test(unittest.TestCase):
21    """Tests for update_kernel_afdo."""
22
23    def make_tempdir(self) -> Path:
24        x = Path(tempfile.mkdtemp(prefix="update_kernel_afdo_test_"))
25        self.addCleanup(shutil.rmtree, x)
26        return x
27
28    def test_kernel_version_parsing(self):
29        self.assertEqual(
30            update_kernel_afdo.KernelVersion.parse("5.10"),
31            update_kernel_afdo.KernelVersion(major=5, minor=10),
32        )
33
34        with self.assertRaisesRegex(ValueError, ".*invalid kernel version.*"):
35            update_kernel_afdo.KernelVersion.parse("5")
36
37    def test_kernel_version_formatting(self):
38        self.assertEqual(
39            str(update_kernel_afdo.KernelVersion(major=5, minor=10)), "5.10"
40        )
41
42    def test_channel_parsing(self):
43        with self.assertRaisesRegex(ValueError, "No such channel.*"):
44            update_kernel_afdo.Channel.parse("not a channel")
45
46        # Ensure these round-trip.
47        for channel in update_kernel_afdo.Channel:
48            self.assertEqual(
49                channel, update_kernel_afdo.Channel.parse(channel.value)
50            )
51
52    @mock.patch.object(subprocess, "run")
53    def test_branch_autodetection(self, subprocess_run):
54        subprocess_run.return_value = subprocess.CompletedProcess(
55            args=[],
56            returncode=0,
57            stdout=textwrap.dedent(
58                """
59                cros/not-a-release-branch
60                cros/release-R121-15699.B
61                cros/release-R122-15753.B
62                cros/release-R123-15786.B
63                cros/also-not-a-release-branch
64                m/main
65                """
66            ),
67        )
68
69        branch_dict = update_kernel_afdo.autodetect_branches(
70            toolchain_utils=self.make_tempdir()
71        )
72
73        self.assertEqual(
74            branch_dict,
75            {
76                update_kernel_afdo.Channel.CANARY: update_kernel_afdo.GitBranch(
77                    remote="cros",
78                    release_number=124,
79                    branch_name="main",
80                ),
81                update_kernel_afdo.Channel.BETA: update_kernel_afdo.GitBranch(
82                    remote="cros",
83                    release_number=123,
84                    branch_name="release-R123-15786.B",
85                ),
86                update_kernel_afdo.Channel.STABLE: update_kernel_afdo.GitBranch(
87                    remote="cros",
88                    release_number=122,
89                    branch_name="release-R122-15753.B",
90                ),
91            },
92        )
93
94    def test_read_update_cfg_file(self):
95        valid_contents = textwrap.dedent(
96            """
97            # some comment
98            # wow
99            AMD_KVERS="1.0 1.1"
100            ARM_KVERS="1.2"
101            AMD_METADATA_FILE="amd/file/path.json" # comment
102            ARM_METADATA_FILE="arm/file/path.json"
103            """
104        )
105        tmpdir = self.make_tempdir()
106        cfg_path = tmpdir / "test.cfg"
107        cfg_path.write_text(valid_contents, encoding="utf-8")
108        cfg = update_kernel_afdo.read_update_cfg_file(tmpdir, cfg_path)
109        expected_amd64 = update_kernel_afdo.ArchUpdateConfig(
110            versions_to_track=[
111                update_kernel_afdo.KernelVersion(1, 0),
112                update_kernel_afdo.KernelVersion(1, 1),
113            ],
114            metadata_file=tmpdir / "amd/file/path.json",
115        )
116        expected_arm = update_kernel_afdo.ArchUpdateConfig(
117            versions_to_track=[
118                update_kernel_afdo.KernelVersion(1, 2),
119            ],
120            metadata_file=tmpdir / "arm/file/path.json",
121        )
122
123        self.assertEqual(
124            cfg,
125            {
126                update_kernel_afdo.Arch.AMD64: expected_amd64,
127                update_kernel_afdo.Arch.ARM: expected_arm,
128            },
129        )
130
131    def test_parse_kernel_gs_profile(self):
132        timestamp = datetime.datetime.fromtimestamp(1234, datetime.timezone.utc)
133        profile = update_kernel_afdo.KernelGsProfile.from_file_name(
134            timestamp,
135            "R124-15808.0-1710149961.gcov.xz",
136        )
137        self.assertEqual(
138            profile,
139            update_kernel_afdo.KernelGsProfile(
140                release_number=124,
141                chrome_build="15808.0",
142                cwp_timestamp=1710149961,
143                suffix=".gcov.xz",
144                gs_timestamp=timestamp,
145            ),
146        )
147
148    def test_kernel_gs_profile_file_name(self):
149        timestamp = datetime.datetime.fromtimestamp(1234, datetime.timezone.utc)
150        profile = update_kernel_afdo.KernelGsProfile.from_file_name(
151            timestamp,
152            "R124-15808.0-1710149961.gcov.xz",
153        )
154        self.assertEqual(profile.file_name_no_suffix, "R124-15808.0-1710149961")
155        self.assertEqual(profile.file_name, "R124-15808.0-1710149961.gcov.xz")
156
157    def test_gs_time_parsing(self):
158        self.assertEqual(
159            update_kernel_afdo.datetime_from_gs_time("2024-03-04T10:38:50Z"),
160            datetime.datetime(
161                year=2024,
162                month=3,
163                day=4,
164                hour=10,
165                minute=38,
166                second=50,
167                tzinfo=datetime.timezone.utc,
168            ),
169        )
170
171    @mock.patch.object(subprocess, "run")
172    def test_kernel_profile_fetcher_works(self, subprocess_run):
173        subprocess_run.return_value = subprocess.CompletedProcess(
174            args=[],
175            returncode=0,
176            # Don't use textwrap.dedent; linter complains about the line being
177            # too long in that case.
178            stdout="""
179753112  2024-03-04T10:38:50Z gs://here/5.4/R124-15786.10-1709548729.gcov.xz
180TOTAL: 2 objects, 1234 bytes (1.1KiB)
181""",
182        )
183
184        fetcher = update_kernel_afdo.KernelProfileFetcher()
185        results = fetcher.fetch("gs://here/5.4")
186
187        expected_results = [
188            update_kernel_afdo.KernelGsProfile.from_file_name(
189                update_kernel_afdo.datetime_from_gs_time(
190                    "2024-03-04T10:38:50Z"
191                ),
192                "R124-15786.10-1709548729.gcov.xz",
193            ),
194        ]
195        self.assertEqual(results, expected_results)
196
197    @mock.patch.object(subprocess, "run")
198    def test_kernel_profile_fetcher_handles_no_profiles(self, subprocess_run):
199        subprocess_run.return_value = subprocess.CompletedProcess(
200            args=[],
201            returncode=1,
202            stderr="\nCommandException: One or more URLs matched no objects.\n",
203        )
204
205        fetcher = update_kernel_afdo.KernelProfileFetcher()
206        results = fetcher.fetch("gs://here/5.4")
207        self.assertEqual(results, [])
208
209    @mock.patch.object(subprocess, "run")
210    def test_kernel_profile_fetcher_caches_urls(self, subprocess_run):
211        subprocess_run.return_value = subprocess.CompletedProcess(
212            args=[],
213            returncode=0,
214            # Don't use textwrap.dedent; linter complains about the line being
215            # too long in that case.
216            stdout="""
217753112  2024-03-04T10:38:50Z gs://here/5.4/R124-15786.10-1709548729.gcov.xz
218TOTAL: 2 objects, 1234 bytes (1.1KiB)
219""",
220        )
221
222        fetcher = update_kernel_afdo.KernelProfileFetcher()
223        # Fetch these twice, and assert both that:
224        # - Only one fetch is performed.
225        # - Mutating the first list won't impact the later fetch.
226        result = fetcher.fetch("gs://here/5.4")
227        self.assertEqual(len(result), 1)
228        del result[:]
229        result = fetcher.fetch("gs://here/5.4")
230        self.assertEqual(len(result), 1)
231        subprocess_run.assert_called_once()
232
233    @mock.patch.object(update_kernel_afdo.KernelProfileFetcher, "fetch")
234    def test_newest_afdo_artifact_finding_works(self, fetch):
235        late = update_kernel_afdo.KernelGsProfile.from_file_name(
236            datetime.datetime.fromtimestamp(1236, datetime.timezone.utc),
237            "R124-15786.10-1709548729.gcov.xz",
238        )
239        early = update_kernel_afdo.KernelGsProfile.from_file_name(
240            datetime.datetime.fromtimestamp(1234, datetime.timezone.utc),
241            "R124-99999.99-9999999999.gcov.xz",
242        )
243        fetch.return_value = [early, late]
244
245        self.assertEqual(
246            update_kernel_afdo.find_newest_afdo_artifact(
247                update_kernel_afdo.KernelProfileFetcher(),
248                update_kernel_afdo.Arch.AMD64,
249                update_kernel_afdo.KernelVersion(5, 4),
250                release_number=124,
251            ),
252            late,
253        )
254
255    def test_afdo_descriptor_file_round_trips(self):
256        tmpdir = self.make_tempdir()
257        file_path = tmpdir / "desc-file.json"
258
259        contents = {
260            update_kernel_afdo.KernelVersion(5, 10): "file1",
261            update_kernel_afdo.KernelVersion(5, 15): "file2",
262        }
263        self.assertTrue(
264            update_kernel_afdo.write_afdo_descriptor_file(file_path, contents)
265        )
266        self.assertEqual(
267            update_kernel_afdo.read_afdo_descriptor_file(file_path),
268            contents,
269        )
270
271    def test_afdo_descriptor_file_refuses_to_rewrite_identical_contents(self):
272        tmpdir = self.make_tempdir()
273        file_path = tmpdir / "desc-file.json"
274
275        contents = {
276            update_kernel_afdo.KernelVersion(5, 10): "file1",
277            update_kernel_afdo.KernelVersion(5, 15): "file2",
278        }
279        self.assertTrue(
280            update_kernel_afdo.write_afdo_descriptor_file(file_path, contents)
281        )
282        self.assertFalse(
283            update_kernel_afdo.write_afdo_descriptor_file(file_path, contents)
284        )
285
286    def test_repo_autodetects_nothing_if_no_repo_dir(self):
287        self.assertIsNone(
288            update_kernel_afdo.find_chromeos_tree_root(
289                Path("/does/not/exist/nor/is/under/a/repo")
290            )
291        )
292
293    def test_repo_autodetects_repo_dir_correctly(self):
294        tmpdir = self.make_tempdir()
295        test_subdir = tmpdir / "a/directory/and/another/one"
296        test_subdir.mkdir(parents=True)
297        (tmpdir / ".repo").mkdir()
298        self.assertEqual(
299            tmpdir, update_kernel_afdo.find_chromeos_tree_root(test_subdir)
300        )
301
302
303if __name__ == "__main__":
304    unittest.main()
305