Skip to main content

xtask/tasks/guest_test/
download_image.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4use crate::Xtask;
5use anyhow::Context;
6use clap::Parser;
7use clap::ValueEnum;
8use std::path::PathBuf;
9use std::process::Command;
10use vmm_test_images::CONTAINER;
11use vmm_test_images::KnownTestArtifacts;
12use vmm_test_images::STORAGE_ACCOUNT;
13
14/// Download an image from Azure Blob Storage.
15///
16/// If no specific images are specified this command will download all available images.
17#[derive(Parser)]
18pub struct DownloadImageTask {
19    /// The folder to download the images to.
20    #[clap(short, long, default_value = "images")]
21    output_folder: PathBuf,
22    /// The test artifacts to download.
23    #[clap(long)]
24    artifacts: Vec<KnownTestArtifacts>,
25    /// Redownload images even if the file already exists.
26    #[clap(short, long)]
27    force: bool,
28}
29
30impl Xtask for DownloadImageTask {
31    fn run(mut self, _ctx: crate::XtaskCtx) -> anyhow::Result<()> {
32        if self.artifacts.is_empty() {
33            self.artifacts = KnownTestArtifacts::value_variants().to_vec();
34        }
35
36        let filenames = self
37            .artifacts
38            .into_iter()
39            .map(|x| x.filename())
40            .collect::<Vec<_>>();
41
42        if !self.output_folder.exists() {
43            std::fs::create_dir(&self.output_folder)?;
44        }
45
46        let vhd_list = filenames.join(";");
47        run_azcopy_command(&[
48            "copy",
49            &format!("https://{STORAGE_ACCOUNT}.blob.core.windows.net/{CONTAINER}/*"),
50            self.output_folder.to_str().unwrap(),
51            "--include-path",
52            &vhd_list,
53            "--overwrite",
54            &self.force.to_string(),
55        ])?;
56
57        Ok(())
58    }
59}
60
61fn run_azcopy_command(args: &[&str]) -> anyhow::Result<Option<String>> {
62    let azcopy_cmd =
63        which::which("azcopy").context("Failed to find `azcopy`. Is AzCopy installed?")?;
64
65    let mut cmd = Command::new(azcopy_cmd);
66    cmd.args(args);
67
68    let mut child = cmd.spawn().context("Failed to run `azcopy` command.")?;
69    let exit = child
70        .wait()
71        .context("Failed to wait for 'azcopy' command.")?;
72    anyhow::ensure!(exit.success(), "azcopy command failed.");
73    Ok(None)
74}