| #!/usr/bin/env vpython3 |
| # Copyright 2026 The Chromium Authors |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| import importlib.util |
| import os |
| import shutil |
| import stat |
| import subprocess |
| import tempfile |
| import unittest |
| from unittest import mock |
| from bs4 import BeautifulSoup |
| import requests |
| |
| # Load the script under test dynamically |
| spec = importlib.util.spec_from_file_location( |
| 'download_extract', |
| os.path.join(os.path.dirname(__file__), 'download_extract.py'), |
| ) |
| de = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(de) |
| |
| |
| class GetArchUnittest(unittest.TestCase): |
| |
| @mock.patch.dict(os.environ, {'_3PP_PLATFORM': 'linux-amd64'}, clear=True) |
| def test_amd64(self): |
| self.assertEqual(de.get_arch(), 'amd64') |
| |
| @mock.patch.dict(os.environ, {'_3PP_PLATFORM': 'linux-arm64'}, clear=True) |
| def test_arm64(self): |
| self.assertEqual(de.get_arch(), 'arm64') |
| |
| @mock.patch.dict(os.environ, {'_3PP_PLATFORM': 'windows-amd64'}, clear=True) |
| def test_unsupported(self): |
| with self.assertRaisesRegex(RuntimeError, |
| 'Unsupported target platform windows-amd64'): |
| de.get_arch() |
| |
| |
| class GetIsoInfoUnittest(unittest.TestCase): |
| |
| def setUp(self): |
| self._get_patcher = mock.patch.object(de.requests, 'get') |
| self._get_mock = self._get_patcher.start() |
| self.addCleanup(self._get_patcher.stop) |
| |
| def _set_response(self, status_code, text): |
| mock_response = mock.Mock() |
| mock_response.status_code = status_code |
| mock_response.text = text |
| mock_response.raise_for_status.side_effect = None |
| if status_code != 200: |
| mock_response.raise_for_status.side_effect = ( |
| requests.exceptions.HTTPError(f'{status_code}')) |
| self._get_mock.return_value = mock_response |
| |
| def test_amd64(self): |
| html_content = """ |
| <html><body> |
| <a href="ubuntu-24.04.5-live-server-amd64.iso">Download</a> |
| </body></html> |
| """ |
| self._set_response(200, html_content) |
| iso_name, iso_url, version = de.get_iso_info('amd64') |
| self.assertEqual(iso_name, 'ubuntu-24.04.5-live-server-amd64.iso') |
| self.assertEqual( |
| iso_url, |
| 'https://releases.ubuntu.com/noble/ubuntu-24.04.5-live-server-amd64.iso', |
| ) |
| self.assertEqual(version, '24.04.5') |
| self._get_mock.assert_called_once_with('https://releases.ubuntu.com/noble/') |
| |
| def test_arm64(self): |
| html_content = """ |
| <html><body> |
| <a href="ubuntu-24.04.5-live-server-arm64.iso">Download</a> |
| </body></html> |
| """ |
| self._set_response(200, html_content) |
| iso_name, iso_url, version = de.get_iso_info('arm64') |
| self.assertEqual(iso_name, 'ubuntu-24.04.5-live-server-arm64.iso') |
| self.assertEqual( |
| iso_url, |
| 'https://cdimage.ubuntu.com/releases/24.04/release/ubuntu-24.04.5-live-server-arm64.iso', |
| ) |
| self.assertEqual(version, '24.04.5') |
| self._get_mock.assert_called_once_with( |
| 'https://cdimage.ubuntu.com/releases/24.04/release/') |
| |
| def test_http_error(self): |
| self._set_response(404, '') |
| with self.assertRaises(requests.exceptions.HTTPError): |
| de.get_iso_info('amd64') |
| |
| def test_link_not_found(self): |
| html_content = """ |
| <html><body> |
| <p>No ISO links here</p> |
| </body></html> |
| """ |
| self._set_response(200, html_content) |
| with self.assertRaisesRegex(RuntimeError, |
| 'Could not find download link for amd64'): |
| de.get_iso_info('amd64') |
| |
| def test_version_not_found(self): |
| html_content = """ |
| <html><body> |
| <a href="ubuntu-24.04-live-server-amd64.iso">Download</a> |
| </body></html> |
| """ |
| self._set_response(200, html_content) |
| with self.assertRaisesRegex(RuntimeError, |
| 'Could not extract version from ISO name'): |
| de.get_iso_info('amd64') |
| |
| def test_unknown_arch(self): |
| with self.assertRaisesRegex(ValueError, 'Unknown arch: foo'): |
| de.get_iso_info('foo') |
| |
| |
| class CmdGetLatestVersionUnittest(unittest.TestCase): |
| |
| @mock.patch('builtins.print') |
| @mock.patch.object( |
| de, 'get_iso_info', return_value=('name', 'url', '24.04.5')) |
| @mock.patch.object(de, 'get_arch', return_value='amd64') |
| def test_success(self, mock_get_arch, mock_get_iso_info, mock_print): |
| de.cmd_get_latest_version() |
| mock_print.assert_called_once_with('24.04.5') |
| |
| |
| class CmdCheckoutUnittest(unittest.TestCase): |
| |
| def setUp(self): |
| self.addCleanup(mock.patch.stopall) |
| self.mock_get_arch = mock.patch.object( |
| de, 'get_arch', return_value='amd64').start() |
| self.mock_get_iso_info = mock.patch.object(de, 'get_iso_info').start() |
| self.mock_requests_get = mock.patch.object(de.requests, 'get').start() |
| self.mock_open = mock.patch('builtins.open').start() |
| self.mock_exists = mock.patch('os.path.exists').start() |
| self.mock_access = mock.patch('os.access').start() |
| self.mock_tempdir = mock.patch.object(de.tempfile, |
| 'TemporaryDirectory').start() |
| self.mock_run = mock.patch.object(de.subprocess, 'run').start() |
| self.mock_run.return_value = mock.Mock(returncode=0, stdout='', stderr='') |
| |
| self.mock_move = mock.patch.object(de.shutil, 'move').start() |
| self.mock_makedirs = mock.patch.object(de.os, 'makedirs').start() |
| |
| self.iso_name = 'ubuntu-24.04.5-live-server-amd64.iso' |
| self.iso_url = 'http://test.url/' + self.iso_name |
| self.mock_get_iso_info.return_value = ( |
| self.iso_name, |
| self.iso_url, |
| '24.04.5', |
| ) |
| |
| self.tmp_dir = '/tmp/testdir' |
| self.mock_tempdir.return_value.__enter__.return_value = self.tmp_dir |
| self.checkout_path = '/fake/checkout' |
| |
| # Mock successful download |
| mock_resp = mock.Mock() |
| mock_resp.raise_for_status.return_value = None |
| mock_resp.iter_content.return_value = [b'chunk1', b'chunk2'] |
| self.mock_requests_get.return_value.__enter__.return_value = mock_resp |
| |
| # Mock file system |
| self.mock_exists.return_value = True |
| self.mock_access.return_value = True |
| self.fake_7z_dir = '/fake/path/to/7z' |
| self.fake_7z_exe = os.path.join(self.fake_7z_dir, '7zz') |
| |
| @mock.patch.dict( |
| os.environ, |
| { |
| '_3PP_VERSION': '24.04.5', |
| '_7z': '/fake/path/to/7z' |
| }, |
| clear=True, |
| ) |
| def test_success(self): |
| de.cmd_checkout(self.checkout_path) |
| |
| self.mock_get_iso_info.assert_called_once_with('amd64') |
| self.mock_requests_get.assert_called_once_with(self.iso_url, stream=True) |
| self.mock_open.assert_called_once_with(self.iso_name, 'wb') |
| self.mock_exists.assert_any_call(self.fake_7z_exe) |
| self.mock_access.assert_called_once_with(self.fake_7z_exe, os.X_OK) |
| |
| expected_run_calls = [ |
| mock.call( |
| [ |
| self.fake_7z_exe, |
| 'x', |
| os.path.abspath(self.iso_name), |
| '-y', |
| '-o' + self.tmp_dir, |
| ], |
| check=True, |
| capture_output=True, |
| text=True, |
| ), |
| ] |
| self.mock_run.assert_has_calls(expected_run_calls) |
| |
| self.mock_makedirs.assert_called_once_with( |
| self.checkout_path, exist_ok=True) |
| |
| expected_moves = [ |
| mock.call( |
| os.path.join(self.tmp_dir, 'casper/vmlinuz'), |
| os.path.join(self.checkout_path, 'vmlinuz'), |
| ), |
| mock.call( |
| os.path.join(self.tmp_dir, 'casper/initrd'), |
| os.path.join(self.checkout_path, 'initrd'), |
| ), |
| mock.call( |
| os.path.join(self.tmp_dir, 'casper/hwe-vmlinuz'), |
| os.path.join(self.checkout_path, 'hwe-vmlinuz'), |
| ), |
| mock.call( |
| os.path.join(self.tmp_dir, 'casper/hwe-initrd'), |
| os.path.join(self.checkout_path, 'hwe-initrd'), |
| ), |
| mock.call(self.iso_name, os.path.join(self.checkout_path, |
| 'ubuntu.iso')), |
| ] |
| self.mock_move.assert_has_calls(expected_moves, any_order=True) |
| |
| @mock.patch.dict( |
| os.environ, |
| { |
| '_3PP_VERSION': '24.04.4', |
| '_7z': '/fake/path/to/7z' |
| }, |
| clear=True, |
| ) |
| def test_version_mismatch(self): |
| with self.assertRaisesRegex( |
| RuntimeError, |
| 'Requested version 24.04.4 does not match latest 24.04.5'): |
| de.cmd_checkout(self.checkout_path) |
| |
| @mock.patch.dict( |
| os.environ, |
| { |
| '_3PP_VERSION': '24.04.5', |
| '_7z': '/fake/path/to/7z' |
| }, |
| clear=True, |
| ) |
| def test_download_fail(self): |
| mock_resp = self.mock_requests_get.return_value.__enter__.return_value |
| mock_resp.raise_for_status.side_effect = requests.exceptions.HTTPError |
| with self.assertRaises(requests.exceptions.HTTPError): |
| de.cmd_checkout(self.checkout_path) |
| |
| @mock.patch.dict( |
| os.environ, |
| { |
| '_3PP_VERSION': '24.04.5', |
| '_7z': '/fake/path/to/7z' |
| }, |
| clear=True, |
| ) |
| def test_7z_missing(self): |
| # Simulate _7z_exe not existing |
| def exists_side_effect(path): |
| if path == self.fake_7z_exe: |
| return False |
| return True |
| |
| self.mock_exists.side_effect = exists_side_effect |
| with self.assertRaisesRegex(RuntimeError, 'not found or is not executable'): |
| de.cmd_checkout(self.checkout_path) |
| self.mock_exists.assert_any_call(self.fake_7z_exe) |
| self.mock_access.assert_not_called() |
| |
| @mock.patch.dict( |
| os.environ, |
| { |
| '_3PP_VERSION': '24.04.5', |
| '_7z': '/fake/path/to/7z' |
| }, |
| clear=True, |
| ) |
| def test_7z_not_executable(self): |
| self.mock_access.return_value = False |
| with self.assertRaisesRegex(RuntimeError, 'not found or is not executable'): |
| de.cmd_checkout(self.checkout_path) |
| self.mock_exists.assert_any_call(self.fake_7z_exe) |
| self.mock_access.assert_called_once_with(self.fake_7z_exe, os.X_OK) |
| |
| @mock.patch.dict( |
| os.environ, |
| { |
| '_3PP_VERSION': '24.04.5', |
| '_7z': '/fake/path/to/7z' |
| }, |
| clear=True, |
| ) |
| def test_7z_run_fail(self): |
| self.mock_run.side_effect = subprocess.CalledProcessError( |
| returncode=2, cmd='7zz', stderr='ERROR') |
| with self.assertRaises(subprocess.CalledProcessError): |
| de.cmd_checkout(self.checkout_path) |
| self.mock_run.assert_called_once() |
| |
| @mock.patch.dict( |
| os.environ, |
| { |
| '_3PP_VERSION': '24.04.5', |
| '_7z': '/fake/path/to/7z' |
| }, |
| clear=True, |
| ) |
| def test_vmlinuz_not_found(self): |
| |
| def exists_side_effect(path): |
| # Simulation: vmlinuz is missing after extraction |
| if path == os.path.join(self.tmp_dir, 'casper/vmlinuz'): |
| return False |
| return True |
| |
| self.mock_exists.side_effect = exists_side_effect |
| with self.assertRaisesRegex(RuntimeError, 'Could not find vmlinuz'): |
| de.cmd_checkout(self.checkout_path) |
| |
| @mock.patch.dict(os.environ, {'_3PP_VERSION': '24.04.5'}, clear=True) |
| def test_7z_env_var_missing(self): |
| with self.assertRaisesRegex(RuntimeError, |
| 'Error: _7z environment variable not set.'): |
| de.cmd_checkout(self.checkout_path) |
| |
| |
| if __name__ == '__main__': |
| unittest.main() |
