| # Copyright 2022 The Chromium Authors |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| from functools import reduce |
| from PB.recipes.infra.windows_image_builder import input as input_pb |
| from PB.recipes.infra.windows_image_builder import actions |
| from PB.recipes.infra.windows_image_builder import vm |
| from PB.recipes.infra.windows_image_builder import drive |
| from PB.recipes.infra.windows_image_builder import sources |
| from PB.recipes.infra.windows_image_builder import windows_vm |
| from PB.recipes.infra.windows_image_builder import windows_iso |
| from PB.recipes.infra.windows_image_builder \ |
| import windows_image_builder as wib |
| from PB.recipes.infra.windows_image_builder \ |
| import offline_winpe_customization as owc |
| from PB.recipes.infra.windows_image_builder \ |
| import online_windows_customization as onwc |
| from PB.go.chromium.org.luci.buildbucket.proto \ |
| import builds_service as bs_pb2 |
| from PB.go.chromium.org.luci.buildbucket.proto \ |
| import build as b_pb2 |
| |
| from recipe_engine.post_process import DropExpectation, StatusFailure |
| from recipe_engine.post_process import StatusSuccess, StepCommandRE |
| from RECIPE_MODULES.infra.windows_scripts_executor import test_helper as t |
| |
| DEPS = [ |
| 'recipe_engine/path', |
| 'recipe_engine/properties', |
| 'recipe_engine/platform', |
| 'recipe_engine/json', |
| 'recipe_engine/raw_io', |
| 'windows_scripts_executor', |
| 'depot_tools/bot_update', |
| 'depot_tools/gclient', |
| 'recipe_engine/context', |
| 'recipe_engine/file', |
| 'recipe_engine/proto', |
| 'recipe_engine/step', |
| 'recipe_engine/buildbucket', |
| 'recipe_engine/runtime', |
| 'windows_adk', |
| ] |
| |
| PYTHON_VERSION_COMPATIBILITY = 'PY3' |
| |
| PROPERTIES = input_pb.Inputs |
| |
| ################################ TEST DATA #################################### |
| |
| ############################################################################### |
| # The tests here ensure the correctness of the scheduler. The customizations # |
| # here are all designed to test one or the other features of the scheduler. # |
| # In a nutshell scheduler ensures that every customization is executed once # |
| # and only once and every customization can be executed when scheduled. The # |
| # scheduling logic groups together, the customizations that target the same # |
| # arch and can that can be executed on the same builder. There can be max of # |
| # MAX_CUST_BATCH_SIZE customizations in any job scheduled. # |
| # # |
| # The following are the images that are being used for the tests that follow # |
| # 1.ONLINE_BOOTSTRAP: This image contains seven customizations, 6 of them are # |
| # offline winpe customizations that can be run on Wim Customzation Builder # |
| # and the other is a windows iso customization. This can only be run on # |
| # Windows Customization Builder. The scheduler should ensure that wim # |
| # customizations are run first. Followed by windows customization. # |
| # 2.ONLINE_CUST: This image uses the bootstrapped ISO image generated by the # |
| # ONLINE_BOOTSTRAP. This can only be run after the ISO image is generated. # |
| # is also possible to package it with the customization that generates the # |
| # ISO image, but it has to come after the image is generated. # |
| ############################################################################### |
| |
| ONLINE_BOOTSTRAP = wib.Image( |
| name='bootstrap', |
| arch=wib.ARCH_AMD64, |
| customizations=[ |
| wib.Customization( |
| offline_winpe_customization=owc.OfflineWinPECustomization( |
| name='wim-1', |
| offline_customization=[ |
| actions.OfflineAction(name='wim-1-custs', actions=[]) |
| ])), |
| wib.Customization( |
| offline_winpe_customization=owc.OfflineWinPECustomization( |
| name='wim-2', |
| # Use the image from wim-1 |
| image_src=sources.Src( |
| local_src='image(bootstrap)-cust(wim-1)-output'), |
| offline_customization=[ |
| actions.OfflineAction(name='wim-2-custs', actions=[]) |
| ])), |
| wib.Customization( |
| offline_winpe_customization=owc.OfflineWinPECustomization( |
| name='wim-3', |
| # Use the image from wim-2 |
| image_src=sources.Src( |
| local_src='image(bootstrap)-cust(wim-2)-output'), |
| offline_customization=[ |
| actions.OfflineAction(name='wim-3-custs', actions=[]) |
| ])), |
| wib.Customization( |
| offline_winpe_customization=owc.OfflineWinPECustomization( |
| name='wim-4', |
| # Use the image from wim-3 |
| image_src=sources.Src( |
| local_src='image(bootstrap)-cust(wim-3)-output'), |
| offline_customization=[ |
| actions.OfflineAction(name='wim-4-custs', actions=[]) |
| ])), |
| wib.Customization( |
| offline_winpe_customization=owc.OfflineWinPECustomization( |
| name='wim-5', |
| # Use the image from wim-4 |
| image_src=sources.Src( |
| local_src='image(bootstrap)-cust(wim-4)-output'), |
| offline_customization=[ |
| actions.OfflineAction(name='wim-5-custs', actions=[]) |
| ])), |
| wib.Customization( |
| offline_winpe_customization=owc.OfflineWinPECustomization( |
| name='wim-6', |
| # Use the image from wim-5 |
| image_src=sources.Src( |
| local_src='image(bootstrap)-cust(wim-5)-output'), |
| offline_customization=[ |
| actions.OfflineAction(name='wim-6-custs', actions=[]) |
| ])), |
| wib.Customization( |
| windows_iso_customization=windows_iso.WinISOImage( |
| name='bootstrap_image', |
| base_image=sources.Src( |
| cipd_src=sources.CIPDSrc( |
| package='infra/labs/win10', |
| refs='latest', |
| platform='windows-amd64', |
| filename='Win10.iso')), |
| copy_files=[ |
| windows_iso.CopyArtifact( |
| artifact=sources.Src( |
| local_src='image(bootstrap)-cust(wim-6)-output'), |
| mount=True, |
| source='sources/boot/boot.wim', |
| ) |
| ], |
| )) |
| ]) |
| |
| ONLINE_CUST = wib.Image( |
| name='online_cust_example', |
| arch=wib.ARCH_AMD64, |
| customizations=[ |
| wib.Customization( |
| online_windows_customization=onwc.OnlineWinCustomization( |
| name='run_cust', |
| online_customizations=[ |
| onwc.OnlineCustomization( |
| name='test_boot1', |
| vm_config=vm.VM( |
| qemu_vm=vm.QEMU_VM( |
| name='squidward', |
| version='latest', |
| smp='cores=8', |
| memory=8192, |
| device=['ide-cd,drive=newWin.iso'], |
| drives=[ |
| drive.Drive( |
| name='Win10.iso', |
| input_src=sources.Src( |
| # Use the image generated by |
| # ONLINE_BOOTSTRAP |
| local_src='image(bootstrap)' |
| '-cust(bootstrap_image)-output'), |
| interface='none', |
| media='cdrom', |
| readonly=True), |
| drive.Drive( |
| name='system.img', |
| interface='none', |
| media='drive', |
| size=1234546, |
| filesystem='fat') |
| ])), |
| ) |
| ])) |
| ]) |
| |
| TESTS = {'bootstrap.cfg': ONLINE_BOOTSTRAP, 'online.cfg': ONLINE_CUST} |
| DIR_DATA = {'tests/basic': ['bootstrap.cfg', 'online.cfg']} |
| |
| |
| def tests(config): |
| return TESTS[config] |
| |
| |
| def lsdir(path): |
| return DIR_DATA[path] |
| |
| |
| ############################## TEST DATA END ################################## |
| |
| |
| def RunSteps(api, config): |
| api.windows_scripts_executor.init() |
| custs = [] |
| images = [] |
| cfg_path = api.path.cache_dir / config.config_path |
| # list the dir to get the corresponding inputs to tests. see DIR_DATA |
| cfgs = api.file.listdir( |
| "Read all the configs", cfg_path, test_data=lsdir(config.config_path)) |
| for cfg in cfgs: |
| # read the configs using helper method tests, see TESTS |
| images.append( |
| api.file.read_proto( |
| name='Reading ' + str(cfg), |
| source=cfg, |
| msg_class=wib.Image, |
| codec='TEXTPB', |
| test_proto=tests(api.path.basename(cfg)))) |
| # Initialize all the images |
| for image in images: |
| custs.extend(api.windows_scripts_executor.init_customizations(image)) |
| |
| # Get all the inputs required. This will be used to determine if we have |
| # to cache any images in online customization |
| inputs = [] |
| for cust in custs: |
| for ip in cust.inputs: |
| if ip.WhichOneof('src') == 'local_src': |
| inputs.append(ip.local_src) |
| |
| # Process all the customizations |
| custs = api.windows_scripts_executor.process_customizations(custs, {}, inputs) |
| |
| configs = api.windows_scripts_executor.get_executable_configs(custs) |
| executed_custs = set() |
| while configs: |
| for builder, images in configs.items(): |
| for image, keys in images: |
| if not keys.issubset(executed_custs): |
| executed_custs.update(keys) |
| api.file.write_proto('Write image', api.path.cache_dir / image.name, |
| image, 'TEXTPB') |
| else: |
| raise Exception('Failed to execute the last run') #pragma no cover |
| configs = api.windows_scripts_executor.get_executable_configs(custs) |
| |
| |
| def GenTests(api): |
| # keys corresponding to the test customizations |
| key_wim_1 = '861f1938f03821f71380f030cb4310d95e66f93ce299cc11e46b65a1405799ba' |
| key_wim_2 = '0268c9590f99bb16340b5850394995d32716d93439152d37bf47918872290e9c' |
| key_wim_3 = 'e1a4cfb3b0533403a04dcab86d20af0b8913bbc5fd9c26044aefb5ac9b45e551' |
| key_wim_4 = 'f8aae81ee074af51989d85af4820d3ec0f1e33af1862c8008954e80c2c7b6237' |
| key_wim_5 = 'faaaaea5ecd871b7885522d379843731efb09b43727bafc6e895fd2c0e592430' |
| key_wim_6 = 'd7ba1e059a49f980e58e704411b7134976336a16e03bf916697483b14af5d356' |
| key_boot = 'bbcff59ef2de5ccb05c8b9679ef78b371d35cc5d58dec639422e34b16bd25219' |
| key_img = 'e4eb670abbd2e02a9abe006c36aa439c7e931740b6ca1cd4ee48ec4ca5eab363' |
| |
| system = 'boot(test_boot1)-drive(system.img)-output.zip' |
| |
| def mock_gsutil(api, url, index): |
| """ mock_wim mocks the existence of output given by the url |
| |
| mock the gsutil check. Mock file doesn't exist for index-1 times and |
| that the file exists for the index time it's queried. |
| |
| Args: |
| url: The gsutil url representing the file |
| index: Index of the query for which we respond yes the file exists. |
| """ |
| return reduce(lambda a, b: a + b, [ |
| t.MOCK_CUST_OUTPUT( |
| api, url + |
| (' ({})'.format(postfix) if postfix != 1 else ''), postfix == index) |
| for postfix in range(1, index + 1) |
| ]) |
| |
| yield ( |
| api.test('nothing scheduled', api.platform('linux', 64)) + |
| # Run the test with basic images |
| api.properties(input_pb.Inputs(config_path="tests/basic")) + |
| api.properties.environ(input_pb.EnvProperties(MAX_CUST_BATCH_SIZE="5")) + |
| # Scheduler gets exists response for all the queries. It doesn't have to |
| # execute anything |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_1), 1) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_2), 1) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_3), 1) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_4), 1) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_5), 1) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_6), 1) + |
| mock_gsutil(api, 'gs://chrome-gce-images/WIB-ISO/{}.iso'.format(key_boot), |
| 1) + api.post_process(StatusSuccess) + |
| api.post_process(DropExpectation)) |
| |
| yield ( |
| api.test('schedule_online_cust', api.platform('linux', 64)) + |
| # Run the test with basic images |
| api.properties(input_pb.Inputs(config_path="tests/basic")) + |
| api.properties.environ(input_pb.EnvProperties(MAX_CUST_BATCH_SIZE="5")) + |
| # Scheduler only needs to schedule online builder |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_1), 1) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_2), 1) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_3), 1) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_4), 1) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_5), 1) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_6), 1) + |
| mock_gsutil(api, 'gs://chrome-gce-images/WIB-ISO/{}.iso'.format(key_boot), |
| 1) + api.post_process(StatusSuccess) + |
| api.post_process(DropExpectation)) |
| |
| yield ( |
| api.test('schedule_iso_cust', api.platform('linux', 64)) + |
| # Run the test with basic images |
| api.properties(input_pb.Inputs(config_path="tests/basic")) + mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_1), 1) + |
| api.properties.environ(input_pb.EnvProperties(MAX_CUST_BATCH_SIZE="5")) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_2), 1) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_3), 1) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_4), 1) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_5), 1) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_6), 1) + |
| mock_gsutil(api, 'gs://chrome-gce-images/WIB-ISO/{}.iso'.format(key_boot), |
| 2) + api.post_process(StatusSuccess) + |
| api.post_process(DropExpectation)) |
| |
| yield ( |
| api.test('schedule_wim_cust', api.platform('linux', 64)) + |
| api.properties.environ(input_pb.EnvProperties(MAX_CUST_BATCH_SIZE="5")) + |
| # Run the test with basic images |
| api.properties(input_pb.Inputs(config_path="tests/basic")) + mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_1), 2) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_2), 2) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_3), 2) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_4), 2) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_5), 3) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_6), 3) + |
| mock_gsutil(api, 'gs://chrome-gce-images/WIB-ISO/{}.iso'.format(key_boot), |
| 1) + api.post_process(StatusSuccess) + |
| api.post_process(DropExpectation)) |
| |
| yield ( |
| api.test('happy path', api.platform('linux', 64)) + |
| api.properties.environ(input_pb.EnvProperties(MAX_CUST_BATCH_SIZE="5")) + |
| # Run the test with basic images |
| api.properties(input_pb.Inputs(config_path="tests/basic")) + mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_1), 2) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_2), 2) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_3), 2) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_4), 2) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_5), 3) + |
| mock_gsutil( |
| api, 'gs://chrome-gce-images/WIB-WIM/{}.zip'.format(key_wim_6), 5) + |
| mock_gsutil(api, 'gs://chrome-gce-images/WIB-ISO/{}.iso'.format(key_boot), |
| 3) + api.post_process(StatusSuccess) + |
| api.post_process(DropExpectation)) |
