<commit_before>//! Tests of the interactive console installer

pub mod mock;

use crate::mock::clitools::{
    self, expect_stderr_ok, expect_stdout_ok, set_current_dist_date, Config, SanitizedOutput,
    Scenario,
};
use crate::mock::{get_path, restore_path};
use lazy_static::lazy_static;
use rustup::utils::raw;
use std::fs;
use std::io::Write;
use std::process::Stdio;
use std::sync::Mutex;

pub fn setup_(complex: bool, f: &dyn Fn(&Config)) {
    let scenario = if complex {
        Scenario::UnavailableRls
    } else {
        Scenario::SimpleV2
    };
    clitools::setup(scenario, &|config| {
        // Lock protects environment variables
        lazy_static! {
            static ref LOCK: Mutex<()> = Mutex::new(());
        }
        let _g = LOCK.lock();

        // An windows these tests mess with the user's PATH. Save
        // and restore them here to keep from trashing things.
        let saved_path = get_path();
        let _g = scopeguard::guard(saved_path, restore_path);

        f(config);
    });
}

pub fn setup(f: &dyn Fn(&Config)) {
    setup_(false, f)
}

fn run_input(config: &Config, args: &[&str], input: &str) -> SanitizedOutput {
    let mut cmd = clitools::cmd(config, args[0], &args[1..]);
    clitools::env(config, &mut cmd);

    cmd.stdin(Stdio::piped());
    cmd.stdout(Stdio::piped());
    cmd.stderr(Stdio::piped());
    let mut child = cmd.spawn().unwrap();

    child
        .stdin
        .as_mut()
        .unwrap()
        .write_all(input.as_bytes())
        .unwrap();
    let out = child.wait_with_output().unwrap();

    SanitizedOutput {
        ok: out.status.success(),
        stdout: String::from_utf8(out.stdout).unwrap(),
        stderr: String::from_utf8(out.stderr).unwrap(),
    }
}

#[test]
fn smoke_test() {
    setup(&|config| {
        let out = run_input(config, &["rustup-init"], "\n\n");
        assert!(out.ok);
    });
}

#[test]
fn update() {
    setup(&|config| {
        run_input(config, &["rustup-init"], "\n\n");
        let out = run_input(config, &["rustup-init"], "\n\n");
        assert!(out.ok, "stdout:\n{}\nstderr:\n{}", out.stdout, out.stderr);
    });
}

// Testing that the right number of blank lines are printed after the
// 'pre-install' message and before the 'post-install' message.
#[test]
fn blank_lines_around_stderr_log_output_install() {
    setup(&|config| {
        let out = run_input(config, &["rustup-init"], "\n\n");

        // During an interactive session, after "Press the Enter
        // key..."  the UI emits a blank line, then there is a blank
        // line that comes from the user pressing enter, then log
        // output on stderr, then an explicit blank line on stdout
        // before printing $toolchain installed
        assert!(out.stdout.contains(
            r"
3) Cancel installation
>

  stable installed - 1.1.0 (hash-stable-1.1.0)

Rust is installed now. Great!
"
        ));
    });
}

#[test]
fn blank_lines_around_stderr_log_output_update() {
    setup(&|config| {
        run_input(config, &["rustup-init"], "\n\n");
        let out = run_input(config, &["rustup-init"], "\n\n");
        println!("-- stdout --\n {}", out.stdout);
        println!("-- stderr --\n {}", out.stderr);

        assert!(out.stdout.contains(
            r"
3) Cancel installation
>

Rust is installed now. Great!
"
        ));
    });
}

#[test]
fn user_says_nope() {
    setup(&|config| {
        let out = run_input(config, &["rustup-init"], "n\n\n");
        assert!(out.ok);
        assert!(!config.cargodir.join("bin").exists());
    });
}

#[test]
fn with_no_modify_path() {
    setup(&|config| {
        let out = run_input(config, &["rustup-init", "--no-modify-path"], "\n\n");
        assert!(out.ok);
        assert!(out
            .stdout
            .contains("This path needs to be in your PATH environment variable"));

        if cfg!(unix) {
            assert!(!config.homedir.join(".profile").exists());
        }
    });
}

#[test]
fn with_no_toolchain() {
    setup(&|config| {
        let out = run_input(config, &["rustup-init", "--default-toolchain=none"], "\n\n");
        assert!(out.ok);

        expect_stdout_ok(config, &["rustup", "show"], "no active toolchain");
    });
}

#[test]
fn with_non_default_toolchain() {
    setup(&|config| {
        let out = run_input(
            config,
            &["rustup-init", "--default-toolchain=nightly"],
            "\n\n",
        );
        assert!(out.ok);

        expect_stdout_ok(config, &["rustup", "show"], "nightly");
    });
}

#[test]
fn with_non_release_channel_non_default_toolchain() {
    setup(&|config| {
        let out = run_input(
            config,
            &["rustup-init", "--default-toolchain=nightly-2015-01-02"],
            "\n\n",
        );
        assert!(out.ok);

        expect_stdout_ok(config, &["rustup", "show"], "nightly");
        expect_stdout_ok(config, &["rustup", "show"], "2015-01-02");
    });
}

#[test]
fn set_nightly_toolchain() {
    setup(&|config| {
        let out = run_input(config, &["rustup-init"], "2\n\nnightly\n\n\n\n\n");
        assert!(out.ok);

        expect_stdout_ok(config, &["rustup", "show"], "nightly");
    });
}

#[test]
fn set_no_modify_path() {
    setup(&|config| {
        let out = run_input(config, &["rustup-init"], "2\n\n\n\nno\n\n\n");
        assert!(out.ok);

        if cfg!(unix) {
            assert!(!config.homedir.join(".profile").exists());
        }
    });
}

#[test]
fn set_nightly_toolchain_and_unset() {
    setup(&|config| {
        let out = run_input(
            config,
            &["rustup-init"],
            "2\n\nnightly\n\n\n2\n\nbeta\n\n\n\n\n",
        );
        assert!(out.ok);

        expect_stdout_ok(config, &["rustup", "show"], "beta");
    });
}

#[test]
fn user_says_nope_after_advanced_install() {
    setup(&|config| {
        let out = run_input(config, &["rustup-init"], "2\n\n\n\n\nn\n\n\n");
        assert!(out.ok);
        assert!(!config.cargodir.join("bin").exists());
    });
}

#[test]
fn install_forces_and_skips_rls() {
    setup_(true, &|config| {
        set_current_dist_date(config, "2015-01-01");

        let out = run_input(
            config,
            &[
                "rustup-init",
                "--profile",
                "complete",
                "--default-toolchain",
                "nightly",
            ],
            "\n\n",
        );
        assert!(out.ok);
        assert!(out
            .stderr
            .contains("warning: Force-skipping unavailable component"));
    });
}

#[test]
fn test_warn_if_complete_profile_is_used() {
    setup(&|config| {
        expect_stderr_ok(
            config,
            &["rustup-init", "-y", "--profile", "complete"],
            "warning: downloading with complete profile",
        );
    });
}

fn create_rustup_sh_metadata(config: &Config) {
    let rustup_dir = config.homedir.join(".rustup");
    fs::create_dir_all(&rustup_dir).unwrap();
    let version_file = rustup_dir.join("rustup-version");
    raw::write_file(&version_file, "").unwrap();
}

#[test]
fn test_prompt_fail_if_rustup_sh_already_installed_reply_nothing() {
    setup(&|config| {
        create_rustup_sh_metadata(&config);
        let out = run_input(config, &["rustup-init"], "\n");
        assert!(!out.ok);
        assert!(out
            .stderr
            .contains("warning: it looks like you have existing rustup.sh metadata"));
        assert!(out
            .stderr
            .contains("error: cannot install while rustup.sh is installed"));
        assert!(out.stdout.contains("Continue? (y/N)"));
    })
}

#[test]
fn test_prompt_fail_if_rustup_sh_already_installed_reply_no() {
    setup(&|config| {
        create_rustup_sh_metadata(&config);
        let out = run_input(config, &["rustup-init"], "no\n");
        assert!(!out.ok);
        assert!(out
            .stderr
            .contains("warning: it looks like you have existing rustup.sh metadata"));
        assert!(out
            .stderr
            .contains("error: cannot install while rustup.sh is installed"));
        assert!(out.stdout.contains("Continue? (y/N)"));
    })
}

#[test]
fn test_prompt_succeed_if_rustup_sh_already_installed_reply_yes() {
    setup(&|config| {
        create_rustup_sh_metadata(&config);
        let out = run_input(config, &["rustup-init"], "yes\n\n\n");
        assert!(out.ok);
        assert!(out
            .stderr
            .contains("warning: it looks like you have existing rustup.sh metadata"));
        assert!(out
            .stderr
            .contains("error: cannot install while rustup.sh is installed"));
        assert!(out.stdout.contains("Continue? (y/N)"));
        assert!(!out.stdout.contains(
            "warning: continuing (because the -y flag is set and the error is ignorable)"
        ))
    })
}

#[test]
fn test_warn_succeed_if_rustup_sh_already_installed_y_flag() {
    setup(&|config| {
        create_rustup_sh_metadata(&config);
        let out = run_input(config, &["rustup-init", "-y"], "");
        assert!(out.ok);
        assert!(out
            .stderr
            .contains("warning: it looks like you have existing rustup.sh metadata"));
        assert!(out
            .stderr
            .contains("error: cannot install while rustup.sh is installed"));
        assert!(out.stderr.contains(
            "warning: continuing (because the -y flag is set and the error is ignorable)"
        ));
        assert!(!out.stdout.contains("Continue? (y/N)"));
    })
}

#[ignore] // Can't test environment variable.
#[test]
fn test_succeed_if_rustup_sh_already_installed_env_var_set() {}
<commit_msg>Created `run_input_with_env`.<commit_after>//! Tests of the interactive console installer

pub mod mock;

use crate::mock::clitools::{
    self, expect_stderr_ok, expect_stdout_ok, set_current_dist_date, Config, SanitizedOutput,
    Scenario,
};
use crate::mock::{get_path, restore_path};
use lazy_static::lazy_static;
use rustup::utils::raw;
use std::fs;
use std::io::Write;
use std::process::Stdio;
use std::sync::Mutex;

pub fn setup_(complex: bool, f: &dyn Fn(&Config)) {
    let scenario = if complex {
        Scenario::UnavailableRls
    } else {
        Scenario::SimpleV2
    };
    clitools::setup(scenario, &|config| {
        // Lock protects environment variables
        lazy_static! {
            static ref LOCK: Mutex<()> = Mutex::new(());
        }
        let _g = LOCK.lock();

        // An windows these tests mess with the user's PATH. Save
        // and restore them here to keep from trashing things.
        let saved_path = get_path();
        let _g = scopeguard::guard(saved_path, restore_path);

        f(config);
    });
}

pub fn setup(f: &dyn Fn(&Config)) {
    setup_(false, f)
}

fn run_input(config: &Config, args: &[&str], input: &str) -> SanitizedOutput {
    run_input_with_env(config, args, input, &[])
}

fn run_input_with_env(
    config: &Config,
    args: &[&str],
    input: &str,
    env: &[(&str, &str)],
) -> SanitizedOutput {
    let mut cmd = clitools::cmd(config, args[0], &args[1..]);
    clitools::env(config, &mut cmd);

    for (key, value) in env.iter() {
        cmd.env(key, value);
    }

    cmd.stdin(Stdio::piped());
    cmd.stdout(Stdio::piped());
    cmd.stderr(Stdio::piped());
    let mut child = cmd.spawn().unwrap();

    child
        .stdin
        .as_mut()
        .unwrap()
        .write_all(input.as_bytes())
        .unwrap();
    let out = child.wait_with_output().unwrap();

    SanitizedOutput {
        ok: out.status.success(),
        stdout: String::from_utf8(out.stdout).unwrap(),
        stderr: String::from_utf8(out.stderr).unwrap(),
    }
}

#[test]
fn smoke_test() {
    setup(&|config| {
        let out = run_input(config, &["rustup-init"], "\n\n");
        assert!(out.ok);
    });
}

#[test]
fn update() {
    setup(&|config| {
        run_input(config, &["rustup-init"], "\n\n");
        let out = run_input(config, &["rustup-init"], "\n\n");
        assert!(out.ok, "stdout:\n{}\nstderr:\n{}", out.stdout, out.stderr);
    });
}

// Testing that the right number of blank lines are printed after the
// 'pre-install' message and before the 'post-install' message.
#[test]
fn blank_lines_around_stderr_log_output_install() {
    setup(&|config| {
        let out = run_input(config, &["rustup-init"], "\n\n");

        // During an interactive session, after "Press the Enter
        // key..."  the UI emits a blank line, then there is a blank
        // line that comes from the user pressing enter, then log
        // output on stderr, then an explicit blank line on stdout
        // before printing $toolchain installed
        assert!(out.stdout.contains(
            r"
3) Cancel installation
>

  stable installed - 1.1.0 (hash-stable-1.1.0)

Rust is installed now. Great!
"
        ));
    });
}

#[test]
fn blank_lines_around_stderr_log_output_update() {
    setup(&|config| {
        run_input(config, &["rustup-init"], "\n\n");
        let out = run_input(config, &["rustup-init"], "\n\n");
        println!("-- stdout --\n {}", out.stdout);
        println!("-- stderr --\n {}", out.stderr);

        assert!(out.stdout.contains(
            r"
3) Cancel installation
>

Rust is installed now. Great!
"
        ));
    });
}

#[test]
fn user_says_nope() {
    setup(&|config| {
        let out = run_input(config, &["rustup-init"], "n\n\n");
        assert!(out.ok);
        assert!(!config.cargodir.join("bin").exists());
    });
}

#[test]
fn with_no_modify_path() {
    setup(&|config| {
        let out = run_input(config, &["rustup-init", "--no-modify-path"], "\n\n");
        assert!(out.ok);
        assert!(out
            .stdout
            .contains("This path needs to be in your PATH environment variable"));

        if cfg!(unix) {
            assert!(!config.homedir.join(".profile").exists());
        }
    });
}

#[test]
fn with_no_toolchain() {
    setup(&|config| {
        let out = run_input(config, &["rustup-init", "--default-toolchain=none"], "\n\n");
        assert!(out.ok);

        expect_stdout_ok(config, &["rustup", "show"], "no active toolchain");
    });
}

#[test]
fn with_non_default_toolchain() {
    setup(&|config| {
        let out = run_input(
            config,
            &["rustup-init", "--default-toolchain=nightly"],
            "\n\n",
        );
        assert!(out.ok);

        expect_stdout_ok(config, &["rustup", "show"], "nightly");
    });
}

#[test]
fn with_non_release_channel_non_default_toolchain() {
    setup(&|config| {
        let out = run_input(
            config,
            &["rustup-init", "--default-toolchain=nightly-2015-01-02"],
            "\n\n",
        );
        assert!(out.ok);

        expect_stdout_ok(config, &["rustup", "show"], "nightly");
        expect_stdout_ok(config, &["rustup", "show"], "2015-01-02");
    });
}

#[test]
fn set_nightly_toolchain() {
    setup(&|config| {
        let out = run_input(config, &["rustup-init"], "2\n\nnightly\n\n\n\n\n");
        assert!(out.ok);

        expect_stdout_ok(config, &["rustup", "show"], "nightly");
    });
}

#[test]
fn set_no_modify_path() {
    setup(&|config| {
        let out = run_input(config, &["rustup-init"], "2\n\n\n\nno\n\n\n");
        assert!(out.ok);

        if cfg!(unix) {
            assert!(!config.homedir.join(".profile").exists());
        }
    });
}

#[test]
fn set_nightly_toolchain_and_unset() {
    setup(&|config| {
        let out = run_input(
            config,
            &["rustup-init"],
            "2\n\nnightly\n\n\n2\n\nbeta\n\n\n\n\n",
        );
        assert!(out.ok);

        expect_stdout_ok(config, &["rustup", "show"], "beta");
    });
}

#[test]
fn user_says_nope_after_advanced_install() {
    setup(&|config| {
        let out = run_input(config, &["rustup-init"], "2\n\n\n\n\nn\n\n\n");
        assert!(out.ok);
        assert!(!config.cargodir.join("bin").exists());
    });
}

#[test]
fn install_forces_and_skips_rls() {
    setup_(true, &|config| {
        set_current_dist_date(config, "2015-01-01");

        let out = run_input(
            config,
            &[
                "rustup-init",
                "--profile",
                "complete",
                "--default-toolchain",
                "nightly",
            ],
            "\n\n",
        );
        assert!(out.ok);
        assert!(out
            .stderr
            .contains("warning: Force-skipping unavailable component"));
    });
}

#[test]
fn test_warn_if_complete_profile_is_used() {
    setup(&|config| {
        expect_stderr_ok(
            config,
            &["rustup-init", "-y", "--profile", "complete"],
            "warning: downloading with complete profile",
        );
    });
}

fn create_rustup_sh_metadata(config: &Config) {
    let rustup_dir = config.homedir.join(".rustup");
    fs::create_dir_all(&rustup_dir).unwrap();
    let version_file = rustup_dir.join("rustup-version");
    raw::write_file(&version_file, "").unwrap();
}

#[test]
fn test_prompt_fail_if_rustup_sh_already_installed_reply_nothing() {
    setup(&|config| {
        create_rustup_sh_metadata(&config);
        let out = run_input(config, &["rustup-init"], "\n");
        assert!(!out.ok);
        assert!(out
            .stderr
            .contains("warning: it looks like you have existing rustup.sh metadata"));
        assert!(out
            .stderr
            .contains("error: cannot install while rustup.sh is installed"));
        assert!(out.stdout.contains("Continue? (y/N)"));
    })
}

#[test]
fn test_prompt_fail_if_rustup_sh_already_installed_reply_no() {
    setup(&|config| {
        create_rustup_sh_metadata(&config);
        let out = run_input(config, &["rustup-init"], "no\n");
        assert!(!out.ok);
        assert!(out
            .stderr
            .contains("warning: it looks like you have existing rustup.sh metadata"));
        assert!(out
            .stderr
            .contains("error: cannot install while rustup.sh is installed"));
        assert!(out.stdout.contains("Continue? (y/N)"));
    })
}

#[test]
fn test_prompt_succeed_if_rustup_sh_already_installed_reply_yes() {
    setup(&|config| {
        create_rustup_sh_metadata(&config);
        let out = run_input(config, &["rustup-init"], "yes\n\n\n");
        assert!(out.ok);
        assert!(out
            .stderr
            .contains("warning: it looks like you have existing rustup.sh metadata"));
        assert!(out
            .stderr
            .contains("error: cannot install while rustup.sh is installed"));
        assert!(out.stdout.contains("Continue? (y/N)"));
        assert!(!out.stdout.contains(
            "warning: continuing (because the -y flag is set and the error is ignorable)"
        ))
    })
}

#[test]
fn test_warn_succeed_if_rustup_sh_already_installed_y_flag() {
    setup(&|config| {
        create_rustup_sh_metadata(&config);
        let out = run_input(config, &["rustup-init", "-y"], "");
        assert!(out.ok);
        assert!(out
            .stderr
            .contains("warning: it looks like you have existing rustup.sh metadata"));
        assert!(out
            .stderr
            .contains("error: cannot install while rustup.sh is installed"));
        assert!(out.stderr.contains(
            "warning: continuing (because the -y flag is set and the error is ignorable)"
        ));
        assert!(!out.stdout.contains("Continue? (y/N)"));
    })
}
<|endoftext|>


<commit_before>use std::{env, fs, io};
use std::path::Path;
use std::process::{Command, ExitStatus, exit};

#[derive(Debug, Clone)]
struct Metric<'a> {
    name: &'a str,
    prepare: &'a str,
    cmd: &'a str,
}

impl Metric<'_> {
    fn outfile(&self, project: &str) -> String {
        [self.name, "_", project, ".json"].join("")
    }
}

fn main() {
    // TODO args lib
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        println!("please provide the target projects directory");
        exit(1);
    }
    let projects_directory = &args[1];

    // to add a new metric to the test suite, simply define it in this list:
    let metrics: Vec<Metric> = vec![
        Metric { name:"parse", prepare: "rm -rf target/", cmd: "dbt parse --no-version-check" },
    ];

    // list out all projects
    let path = Path::new(&projects_directory);
    let project_entries = fs::read_dir(path).unwrap();

    let results: Result<Vec<ExitStatus>, io::Error > = project_entries.map(|entry| {
        metrics.clone().into_iter().map(|metric| {
            let path = entry.as_ref().unwrap().path();
            let project_name = &path.file_name().and_then(|x| x.to_str()).unwrap();

            Command::new("hyperfine")
                .current_dir(&path)
                // warms filesystem caches by running the command first without counting it.
                // alternatively we could clear them before each run
                .arg("--warmup 1")
                .arg("--prepare")
                .arg(metric.prepare)
                .arg([metric.cmd, " --profiles-dir ", &projects_directory, "/../project_config/"].join(""))
                .arg("--export-json")
                .arg([&projects_directory, "/", &metric.outfile(project_name)].join(""))
                // this prevents capture dbt output. Noisy, but good for debugging when tests fail.
                .arg("--show-output")
                .status() // use spawn() here instead for more information
        }).collect::<Vec<Result<ExitStatus, io::Error>>>()
    }).flatten().collect();

    // only exit with status code 0 if everything ran as expected
    match results {
        // if dispatch of any of the commands failed, panic with that error
        Err(e) => panic!("{}", e),
        Ok(statuses) => {
            for status in statuses {
                match status.code() {
                    None => (),
                    Some(0) => (),
                    // if any child command exited with a non zero status code, exit with the same one here.
                    Some(nonzero) => exit(nonzero),
                }
            }
            // everything ran as expected
            exit(0);
        },
    }
}
<commit_msg>fix warmup syntax<commit_after>use std::{env, fs, io};
use std::path::Path;
use std::process::{Command, ExitStatus, exit};

#[derive(Debug, Clone)]
struct Metric<'a> {
    name: &'a str,
    prepare: &'a str,
    cmd: &'a str,
}

impl Metric<'_> {
    fn outfile(&self, project: &str) -> String {
        [self.name, "_", project, ".json"].join("")
    }
}

fn main() {
    // TODO args lib
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        println!("please provide the target projects directory");
        exit(1);
    }
    let projects_directory = &args[1];

    // to add a new metric to the test suite, simply define it in this list:
    let metrics: Vec<Metric> = vec![
        Metric { name:"parse", prepare: "rm -rf target/", cmd: "dbt parse --no-version-check" },
    ];

    // list out all projects
    let path = Path::new(&projects_directory);
    let project_entries = fs::read_dir(path).unwrap();

    let results: Result<Vec<ExitStatus>, io::Error > = project_entries.map(|entry| {
        metrics.clone().into_iter().map(|metric| {
            let path = entry.as_ref().unwrap().path();
            let project_name = &path.file_name().and_then(|x| x.to_str()).unwrap();

            Command::new("hyperfine")
                .current_dir(&path)
                // warms filesystem caches by running the command first without counting it.
                // alternatively we could clear them before each run
                .arg("--warmup")
                .arg("1")
                .arg("--prepare")
                .arg(metric.prepare)
                .arg([metric.cmd, " --profiles-dir ", &projects_directory, "/../project_config/"].join(""))
                .arg("--export-json")
                .arg([&projects_directory, "/", &metric.outfile(project_name)].join(""))
                // this prevents capture dbt output. Noisy, but good for debugging when tests fail.
                .arg("--show-output")
                .status() // use spawn() here instead for more information
        }).collect::<Vec<Result<ExitStatus, io::Error>>>()
    }).flatten().collect();

    // only exit with status code 0 if everything ran as expected
    match results {
        // if dispatch of any of the commands failed, panic with that error
        Err(e) => panic!("{:?}", e),
        Ok(statuses) => {
            for status in statuses {
                match status.code() {
                    None => (),
                    Some(0) => (),
                    // if any child command exited with a non zero status code, exit with the same one here.
                    Some(nonzero) => exit(nonzero),
                }
            }
            // everything ran as expected
            exit(0);
        },
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use lazy_static::lazy_static;
use log::info;
use proc_macro2::Group;
use proc_macro2::TokenStream as TokenStream2;
use proc_macro2::TokenTree as TokenTree2;
use quote::{quote, TokenStreamExt};
use std::collections::HashMap;

lazy_static! {
    /// Substitutions from the names constructed by bindgen into those
    /// which cxx uses.
    static ref IDENT_REPLACEMENTS: HashMap<&'static str, String> = {
        let mut map = HashMap::new();
        map.insert("std_unique_ptr", "UniquePtr".to_string());
        map.insert("std_string", "CxxString".to_string());
        map
    };
}

#[derive(Debug)]
pub enum ConvertError {
    NoModFound,
    NoFfiFound,
    NoActualGroupFound,
    NoCFound,
    NoExternCGroupFound,
    ExtraFoundInExtern,
    NoGroupAfterAttribute,
}

enum OuterState {
    Start,
    FoundMod,
    FoundFfi,
}

enum ModFfiState {
    Start,
    FoundAttribute,
    FoundExtern,
    FoundC,
}

enum ExternCState {
    Start,
    FoundAttribute,
}

const OUTPUT_TOKENS: bool = false;

/// Converts the bindings generated by bindgen into a form suitable
/// for use with `cxx`.
/// Tasks current performed:
/// * Replaces certain identifiers e.g. `std_unique_ptr` with `UniquePtr`
/// * Replaces pointers with references
/// * Removes repr attributes
/// * Removes link_name attributes
/// * Adds include! directives
/// * Adds #[cxx::bridge]
/// At the moment this is almost certainly not using the best practice for parsing
/// stuff. It's multiple simple-but-yucky state machines. Can undoubtedly be
/// simplified and made less error-prone: TODO.
pub(crate) struct BridgeConverter {
    include_list: Vec<String>,
}

impl BridgeConverter {
    pub fn new(include_list: Vec<String>) -> Self {
        Self { include_list }
    }

    /// Convert a TokenStream of bindgen-generated bindings to a form
    /// suitable for cxx.
    pub(crate) fn convert(&self, bindings: TokenStream2) -> Result<TokenStream2, ConvertError> {
        let mut state = OuterState::Start;
        let mut new_ts = quote! {
            #[cxx::bridge]
            mod ffi
        };
        for t in bindings {
            state = match state {
                OuterState::Start => match t {
                    TokenTree2::Ident(i) if i.to_string() == "mod" => OuterState::FoundMod,
                    _ => return Err(ConvertError::NoModFound),
                },
                OuterState::FoundMod => match t {
                    TokenTree2::Ident(i) if i.to_string() == "ffi" => OuterState::FoundFfi,
                    _ => return Err(ConvertError::NoFfiFound),
                },
                OuterState::FoundFfi => match t {
                    TokenTree2::Group(g) => {
                        new_ts.append(self.convert_mod_ffi(g)?);
                        OuterState::Start
                    }
                    _ => return Err(ConvertError::NoActualGroupFound),
                },
            }
        }
        Ok(new_ts)
    }

    fn convert_mod_ffi(&self, mod_ffi_group: Group) -> Result<Group, ConvertError> {
        let mut state = ModFfiState::Start;
        let mut new_ts = TokenStream2::new();
        for t in mod_ffi_group.stream() {
            state = match state {
                ModFfiState::Start => match t {
                    TokenTree2::Ident(ref i) if i.to_string() == "extern" => {
                        new_ts.append(t);
                        ModFfiState::FoundExtern
                    }
                    TokenTree2::Punct(p) if p.to_string() == "#" => ModFfiState::FoundAttribute,
                    _ => {
                        new_ts.append(t);
                        ModFfiState::Start
                    }
                },
                ModFfiState::FoundExtern => match t {
                    TokenTree2::Literal(ref l) if l.to_string() == "\"C\"" => {
                        new_ts.append(t);
                        ModFfiState::FoundC
                    }
                    _ => return Err(ConvertError::NoCFound),
                },
                ModFfiState::FoundC => match t {
                    TokenTree2::Group(g) => {
                        new_ts.append(self.convert_extern_c(g)?);
                        ModFfiState::Start
                    }
                    _ => return Err(ConvertError::NoActualGroupFound),
                },
                ModFfiState::FoundAttribute => match t {
                    TokenTree2::Group(_) => ModFfiState::Start,
                    _ => return Err(ConvertError::NoGroupAfterAttribute),
                },
            };
        }
        Ok(Group::new(mod_ffi_group.delimiter(), new_ts))
    }

    fn convert_extern_c(&self, extern_c_group: Group) -> Result<Group, ConvertError> {
        let mut state = ExternCState::Start;
        let mut filtered_ts = TokenStream2::new();
        for t in extern_c_group.stream() {
            state = match state {
                ExternCState::Start => match t {
                    TokenTree2::Punct(p) if p.to_string() == "#" => ExternCState::FoundAttribute,
                    _ => {
                        filtered_ts.append(t);
                        ExternCState::Start
                    }
                },
                ExternCState::FoundAttribute => match t {
                    TokenTree2::Group(_) => ExternCState::Start,
                    _ => return Err(ConvertError::NoGroupAfterAttribute),
                },
            }
        }
        let mut new_ts = TokenStream2::new();
        for incl in &self.include_list {
            new_ts.append_all(quote! {
                include!(#incl);
            });
        }
        new_ts.append_all(self.substitute_idents_and_pointers(filtered_ts, 0));
        Ok(Group::new(extern_c_group.delimiter(), new_ts))
    }

    fn substitute_idents_and_pointers(
        &self,
        bindings: TokenStream2,
        indent: usize,
    ) -> TokenStream2 {
        let prefix = std::iter::repeat(" ").take(indent).collect::<String>();
        let mut new_ts = TokenStream2::new();
        for t in bindings {
            if OUTPUT_TOKENS {
                info!("{}TOKEN: {}", prefix, t);
            }
            let replacement = match t {
                TokenTree2::Ident(i) if i.to_string() == "const" => None,
                TokenTree2::Ident(i) => {
                    let name = i.to_string();
                    let e = IDENT_REPLACEMENTS.get(name.as_str());
                    Some(match e {
                        Some(s) => TokenTree2::Ident(proc_macro2::Ident::new(s, i.span())),
                        None => TokenTree2::Ident(i),
                    })
                }
                TokenTree2::Punct(p) if p.as_char() == '*' => {
                    Some(TokenTree2::Punct(proc_macro2::Punct::new('&', p.spacing())))
                }
                TokenTree2::Punct(p) => Some(TokenTree2::Punct(p)),
                TokenTree2::Literal(l) => Some(TokenTree2::Literal(l)),
                TokenTree2::Group(g) => {
                    let delim = g.delimiter();
                    let replacement_tokens =
                        self.substitute_idents_and_pointers(g.stream(), indent + 1);
                    Some(TokenTree2::Group(proc_macro2::Group::new(
                        delim,
                        replacement_tokens,
                    )))
                }
            };
            if let Some(repl) = replacement {
                new_ts.append(repl);
            }
        }
        new_ts
    }
}
<commit_msg>Add comment.<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use lazy_static::lazy_static;
use log::info;
use proc_macro2::Group;
use proc_macro2::TokenStream as TokenStream2;
use proc_macro2::TokenTree as TokenTree2;
use quote::{quote, TokenStreamExt};
use std::collections::HashMap;

lazy_static! {
    /// Substitutions from the names constructed by bindgen into those
    /// which cxx uses.
    static ref IDENT_REPLACEMENTS: HashMap<&'static str, String> = {
        let mut map = HashMap::new();
        map.insert("std_unique_ptr", "UniquePtr".to_string());
        map.insert("std_string", "CxxString".to_string());
        map
    };
}

#[derive(Debug)]
pub enum ConvertError {
    NoModFound,
    NoFfiFound,
    NoActualGroupFound,
    NoCFound,
    NoExternCGroupFound,
    ExtraFoundInExtern,
    NoGroupAfterAttribute,
}

enum OuterState {
    Start,
    FoundMod,
    FoundFfi,
}

enum ModFfiState {
    Start,
    FoundAttribute,
    FoundExtern,
    FoundC,
}

enum ExternCState {
    Start,
    FoundAttribute,
}

const OUTPUT_TOKENS: bool = false;

/// Converts the bindings generated by bindgen into a form suitable
/// for use with `cxx`.
/// Tasks current performed:
/// * Replaces certain identifiers e.g. `std_unique_ptr` with `UniquePtr`
/// * Replaces pointers with references
/// * Removes repr attributes
/// * Removes link_name attributes
/// * Adds include! directives
/// * Adds #[cxx::bridge]
/// At the moment this is almost certainly not using the best practice for parsing
/// stuff. It's multiple simple-but-yucky state machines. Can undoubtedly be
/// simplified and made less error-prone: TODO. Probably the right thing to do
/// is just manipulate the syn types directly.
pub(crate) struct BridgeConverter {
    include_list: Vec<String>,
}

impl BridgeConverter {
    pub fn new(include_list: Vec<String>) -> Self {
        Self { include_list }
    }

    /// Convert a TokenStream of bindgen-generated bindings to a form
    /// suitable for cxx.
    pub(crate) fn convert(&self, bindings: TokenStream2) -> Result<TokenStream2, ConvertError> {
        let mut state = OuterState::Start;
        let mut new_ts = quote! {
            #[cxx::bridge]
            mod ffi
        };
        for t in bindings {
            state = match state {
                OuterState::Start => match t {
                    TokenTree2::Ident(i) if i.to_string() == "mod" => OuterState::FoundMod,
                    _ => return Err(ConvertError::NoModFound),
                },
                OuterState::FoundMod => match t {
                    TokenTree2::Ident(i) if i.to_string() == "ffi" => OuterState::FoundFfi,
                    _ => return Err(ConvertError::NoFfiFound),
                },
                OuterState::FoundFfi => match t {
                    TokenTree2::Group(g) => {
                        new_ts.append(self.convert_mod_ffi(g)?);
                        OuterState::Start
                    }
                    _ => return Err(ConvertError::NoActualGroupFound),
                },
            }
        }
        Ok(new_ts)
    }

    fn convert_mod_ffi(&self, mod_ffi_group: Group) -> Result<Group, ConvertError> {
        let mut state = ModFfiState::Start;
        let mut new_ts = TokenStream2::new();
        for t in mod_ffi_group.stream() {
            state = match state {
                ModFfiState::Start => match t {
                    TokenTree2::Ident(ref i) if i.to_string() == "extern" => {
                        new_ts.append(t);
                        ModFfiState::FoundExtern
                    }
                    TokenTree2::Punct(p) if p.to_string() == "#" => ModFfiState::FoundAttribute,
                    _ => {
                        new_ts.append(t);
                        ModFfiState::Start
                    }
                },
                ModFfiState::FoundExtern => match t {
                    TokenTree2::Literal(ref l) if l.to_string() == "\"C\"" => {
                        new_ts.append(t);
                        ModFfiState::FoundC
                    }
                    _ => return Err(ConvertError::NoCFound),
                },
                ModFfiState::FoundC => match t {
                    TokenTree2::Group(g) => {
                        new_ts.append(self.convert_extern_c(g)?);
                        ModFfiState::Start
                    }
                    _ => return Err(ConvertError::NoActualGroupFound),
                },
                ModFfiState::FoundAttribute => match t {
                    TokenTree2::Group(_) => ModFfiState::Start,
                    _ => return Err(ConvertError::NoGroupAfterAttribute),
                },
            };
        }
        Ok(Group::new(mod_ffi_group.delimiter(), new_ts))
    }

    fn convert_extern_c(&self, extern_c_group: Group) -> Result<Group, ConvertError> {
        let mut state = ExternCState::Start;
        let mut filtered_ts = TokenStream2::new();
        for t in extern_c_group.stream() {
            state = match state {
                ExternCState::Start => match t {
                    TokenTree2::Punct(p) if p.to_string() == "#" => ExternCState::FoundAttribute,
                    _ => {
                        filtered_ts.append(t);
                        ExternCState::Start
                    }
                },
                ExternCState::FoundAttribute => match t {
                    TokenTree2::Group(_) => ExternCState::Start,
                    _ => return Err(ConvertError::NoGroupAfterAttribute),
                },
            }
        }
        let mut new_ts = TokenStream2::new();
        for incl in &self.include_list {
            new_ts.append_all(quote! {
                include!(#incl);
            });
        }
        new_ts.append_all(self.substitute_idents_and_pointers(filtered_ts, 0));
        Ok(Group::new(extern_c_group.delimiter(), new_ts))
    }

    fn substitute_idents_and_pointers(
        &self,
        bindings: TokenStream2,
        indent: usize,
    ) -> TokenStream2 {
        let prefix = std::iter::repeat(" ").take(indent).collect::<String>();
        let mut new_ts = TokenStream2::new();
        for t in bindings {
            if OUTPUT_TOKENS {
                info!("{}TOKEN: {}", prefix, t);
            }
            let replacement = match t {
                TokenTree2::Ident(i) if i.to_string() == "const" => None,
                TokenTree2::Ident(i) => {
                    let name = i.to_string();
                    let e = IDENT_REPLACEMENTS.get(name.as_str());
                    Some(match e {
                        Some(s) => TokenTree2::Ident(proc_macro2::Ident::new(s, i.span())),
                        None => TokenTree2::Ident(i),
                    })
                }
                TokenTree2::Punct(p) if p.as_char() == '*' => {
                    Some(TokenTree2::Punct(proc_macro2::Punct::new('&', p.spacing())))
                }
                TokenTree2::Punct(p) => Some(TokenTree2::Punct(p)),
                TokenTree2::Literal(l) => Some(TokenTree2::Literal(l)),
                TokenTree2::Group(g) => {
                    let delim = g.delimiter();
                    let replacement_tokens =
                        self.substitute_idents_and_pointers(g.stream(), indent + 1);
                    Some(TokenTree2::Group(proc_macro2::Group::new(
                        delim,
                        replacement_tokens,
                    )))
                }
            };
            if let Some(repl) = replacement {
                new_ts.append(repl);
            }
        }
        new_ts
    }
}
<|endoftext|>


<commit_before>use Module;
use socket::Socket;

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Gain {
    LV,
    HV,
}

impl ::std::convert::Into<String> for Gain {
    fn into(self) -> String {
        let s = match self {
            Gain::LV => "LV",
            Gain::HV => "HV",
        };

        String::from(s)
    }
}

impl ::std::str::FromStr for Gain {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "LV" => Ok(Gain::LV),
            "HV" => Ok(Gain::HV),
            gain => Err(format!("Unknow gain '{}'", gain)),
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Source {
    IN1,
    IN2,
}

impl ::std::convert::Into<String> for Source {
    fn into(self) -> String {
        let s = match self {
            Source::IN1 => "SOUR1",
            Source::IN2 => "SOUR2",
        };

        String::from(s)
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Decimation {
    DEC_1,
    DEC_8,
    DEC_64,
    DEC_1024,
    DEC_8192,
    DEC_65536,
}

impl ::std::convert::Into<String> for Decimation {
    fn into(self) -> String {
        let s = match self {
            Decimation::DEC_1 => "1",
            Decimation::DEC_8 => "8",
            Decimation::DEC_64 => "64",
            Decimation::DEC_1024 => "1024",
            Decimation::DEC_8192 => "8192",
            Decimation::DEC_65536 => "65536",
        };

        String::from(s)
    }
}

impl ::std::str::FromStr for Decimation {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "1" => Ok(Decimation::DEC_1),
            "8" => Ok(Decimation::DEC_8),
            "64" => Ok(Decimation::DEC_64),
            "1024" => Ok(Decimation::DEC_1024),
            "8192" => Ok(Decimation::DEC_8192),
            "65536" => Ok(Decimation::DEC_65536),
            decimation => Err(format!("Unknow decimation '{}'", decimation)),
        }
    }
}

impl ::std::convert::Into<SamplingRate> for Decimation {
    fn into(self) -> SamplingRate {
        match self {
            Decimation::DEC_1 => SamplingRate::RATE_125MHz,
            Decimation::DEC_8 => SamplingRate::RATE_15_6MHz,
            Decimation::DEC_64 => SamplingRate::RATE_1_9MHz,
            Decimation::DEC_1024 => SamplingRate::RATE_103_8kHz,
            Decimation::DEC_8192 => SamplingRate::RATE_15_2kHz,
            Decimation::DEC_65536 => SamplingRate::RATE_1_9kHz,
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum SamplingRate {
    RATE_125MHz,
    RATE_15_6MHz,
    RATE_1_9MHz,
    RATE_103_8kHz,
    RATE_15_2kHz,
    RATE_1_9kHz,
}

impl SamplingRate {
    pub fn get_buffer_duration(self) -> ::std::time::Duration {
        let (s, ns) = match self {
            SamplingRate::RATE_125MHz => (0, 131_072),
            SamplingRate::RATE_15_6MHz => (0, 1_049_000),
            SamplingRate::RATE_1_9MHz => (0, 8_389_000),
            SamplingRate::RATE_103_8kHz => (0, 134_218_000),
            SamplingRate::RATE_15_2kHz => (1, 740_000_000),
            SamplingRate::RATE_1_9kHz => (8, 590_000_000),
        };

        ::std::time::Duration::new(s, ns)
    }
}

impl ::std::fmt::Display for SamplingRate {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let display = match self {
            &SamplingRate::RATE_125MHz => "125 MHz",
            &SamplingRate::RATE_15_6MHz => "15.6 MHz",
            &SamplingRate::RATE_1_9MHz => "1.9 MHz",
            &SamplingRate::RATE_103_8kHz => "103.8 kHz",
            &SamplingRate::RATE_15_2kHz => "15.2 kHz",
            &SamplingRate::RATE_1_9kHz => "1.9 kHz",
        };

        write!(f, "{}", display)
    }
}

impl ::std::convert::Into<String> for SamplingRate {
    fn into(self) -> String {
        let s = match self {
            SamplingRate::RATE_125MHz => "125MHz",
            SamplingRate::RATE_15_6MHz => "15_6MHz",
            SamplingRate::RATE_1_9MHz => "1_9MHz",
            SamplingRate::RATE_103_8kHz => "103_8kHz",
            SamplingRate::RATE_15_2kHz => "15_2kHz",
            SamplingRate::RATE_1_9kHz => "1_9kHz",
        };

        String::from(s)
    }
}

impl ::std::str::FromStr for SamplingRate {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "125000000 Hz" => Ok(SamplingRate::RATE_125MHz),
            "15600000 Hz" => Ok(SamplingRate::RATE_15_6MHz),
            "1900000 Hz" => Ok(SamplingRate::RATE_1_9MHz),
            "103800 Hz" => Ok(SamplingRate::RATE_103_8kHz),
            "15200 Hz" => Ok(SamplingRate::RATE_15_2kHz),
            "1900 Hz" => Ok(SamplingRate::RATE_1_9kHz),
            rate => Err(format!("Unknow sampling rate {}", rate)),
        }
    }
}

#[derive(Clone)]
pub struct Acquire {
    socket: ::std::cell::RefCell<Socket>,
    started: bool,
}

impl ::Module for Acquire {
    fn get_socket<'a>(&'a self) -> ::std::cell::RefMut<'a, ::socket::Socket> {
        self.socket.borrow_mut()
    }
}

impl Acquire {
    pub fn new(socket: Socket) -> Self {
        Acquire {
            socket: ::std::cell::RefCell::new(socket),
            started: false,
        }
    }

    /**
     * Starts acquisition.
     */
    pub fn start(&mut self) {
        self.send("ACQ:START");
        self.started = true;
    }

    /**
     * Stops acquisition.
     */
    pub fn stop(&mut self) {
        self.send("ACQ:STOP");
        self.started = false;
    }

    pub fn is_started(&self) -> bool {
        self.started
    }

    /**
     * Stops acquisition and sets all parameters to default values.
     */
    pub fn reset(&self) {
        self.send("ACQ:RST");
    }

    /**
     * Set decimation factor.
     */
    pub fn set_decimation(&self, decimation: Decimation) {
        self.send(format!("ACQ:DEC {}", Into::<String>::into(decimation)));
    }

    /**
     * Get decimation factor.
     */
    pub fn get_decimation(&self) -> Result<Decimation, String> {
        self.send("ACQ:DEC?");

        self.receive()
            .parse()
    }

    /**
     * Get sampling rate.
     */
    pub fn get_sampling_rate(&self) -> Result<SamplingRate, String> {
        self.send("ACQ:SRAT?");

        self.receive()
            .parse()
    }

    /**
     * Enable averaging.
     */
    pub fn enable_average(&self) {
        self.send("ACQ:AVG ON");
    }

    /**
     * Disable averaging.
     */
    pub fn disable_average(&self) {
        self.send("ACQ:AVG OFF");
    }

    /**
     * Get averaging status.
     */
    pub fn is_average_enabled(&self) -> bool {
        self.send("ACQ:AVG?");

        let message = self.receive();

        match message.as_str() {
            "ON" => true,
            _ => false,
        }
    }

    /**
     * Set gain settings to HIGH or LOW.
     *
     * This gain is referring to jumper settings on Red Pitaya fast analog inputs.
     */
    pub fn set_gain(&self, source: Source, gain: Gain) {
        self.send(format!("ACQ:{}:GAIN {}", Into::<String>::into(source), Into::<String>::into(gain)));
    }

    /**
     * Get gain settings to HIGH or LOW.
     */
    pub fn get_gain(&self, source: Source) -> Result<Gain, String> {
        self.send(format!("ACQ:{}:GAIN?", Into::<String>::into(source)));

        self.receive()
            .parse()
    }
}

#[cfg(test)]
mod test {
    macro_rules! acquire_assert {
        ($f:ident, $e:expr) => {
            let (rx, acquire) = create_acquire();

            acquire.$f();
            assert_eq!($e, rx.recv().unwrap());
        }
    }

    #[test]
    fn test_sampling_rate_get_buffer_duration() {
        let duration = ::std::time::Duration::new(8, 590_000_000);

        assert_eq!(duration, ::acquire::SamplingRate::RATE_1_9kHz.get_buffer_duration());
    }

    #[test]
    fn test_start() {
        let (rx, mut acquire) = create_acquire();

        acquire.start();
        assert_eq!("ACQ:START\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_stop() {
        let (rx, mut acquire) = create_acquire();

        acquire.stop();
        assert_eq!("ACQ:STOP\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_is_started() {
        let (_, mut acquire) = create_acquire();

        assert_eq!(acquire.is_started(), false);
        acquire.start();
        assert_eq!(acquire.is_started(), true);
        acquire.stop();
        assert_eq!(acquire.is_started(), false);
    }

    #[test]
    fn test_reset() {
        acquire_assert!(reset, "ACQ:RST\r\n");
    }

    #[test]
    fn test_set_decimation() {
        let (rx, acquire) = create_acquire();

        acquire.set_decimation(::acquire::Decimation::DEC_8);
        assert_eq!("ACQ:DEC 8\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_decimation() {
        let (_, acquire) = create_acquire();

        assert_eq!(acquire.get_decimation(), Ok(::acquire::Decimation::DEC_1));
    }

    #[test]
    fn test_get_sampling_rate() {
        let (_, acquire) = create_acquire();

        assert_eq!(acquire.get_sampling_rate(), Ok(::acquire::SamplingRate::RATE_125MHz));
    }

    #[test]
    fn test_enable_average() {
        acquire_assert!(enable_average, "ACQ:AVG ON\r\n");
    }

    #[test]
    fn test_disable_average() {
        acquire_assert!(disable_average, "ACQ:AVG OFF\r\n");
    }

    #[test]
    fn test_is_average_enabled() {
        let (_, acquire) = create_acquire();

        assert_eq!(acquire.is_average_enabled(), true);
    }

    #[test]
    fn test_set_gain() {
        let (rx, acquire) = create_acquire();

        acquire.set_gain(::acquire::Source::IN1, ::acquire::Gain::LV);
        assert_eq!("ACQ:SOUR1:GAIN LV\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_gain() {
        let (_, acquire) = create_acquire();

        assert_eq!(acquire.get_gain(::acquire::Source::IN1), Ok(::acquire::Gain::HV));
    }

    fn create_acquire() -> (::std::sync::mpsc::Receiver<String>, ::acquire::Acquire) {
        let (addr, rx) = ::test::launch_server();
        let socket = ::socket::Socket::new(addr);

        (rx, ::acquire::Acquire::new(socket))
    }
}
<commit_msg>[acquire] Add warning about acquire.get_sampling_rate()<commit_after>use Module;
use socket::Socket;

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Gain {
    LV,
    HV,
}

impl ::std::convert::Into<String> for Gain {
    fn into(self) -> String {
        let s = match self {
            Gain::LV => "LV",
            Gain::HV => "HV",
        };

        String::from(s)
    }
}

impl ::std::str::FromStr for Gain {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "LV" => Ok(Gain::LV),
            "HV" => Ok(Gain::HV),
            gain => Err(format!("Unknow gain '{}'", gain)),
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Source {
    IN1,
    IN2,
}

impl ::std::convert::Into<String> for Source {
    fn into(self) -> String {
        let s = match self {
            Source::IN1 => "SOUR1",
            Source::IN2 => "SOUR2",
        };

        String::from(s)
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Decimation {
    DEC_1,
    DEC_8,
    DEC_64,
    DEC_1024,
    DEC_8192,
    DEC_65536,
}

impl ::std::convert::Into<String> for Decimation {
    fn into(self) -> String {
        let s = match self {
            Decimation::DEC_1 => "1",
            Decimation::DEC_8 => "8",
            Decimation::DEC_64 => "64",
            Decimation::DEC_1024 => "1024",
            Decimation::DEC_8192 => "8192",
            Decimation::DEC_65536 => "65536",
        };

        String::from(s)
    }
}

impl ::std::str::FromStr for Decimation {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "1" => Ok(Decimation::DEC_1),
            "8" => Ok(Decimation::DEC_8),
            "64" => Ok(Decimation::DEC_64),
            "1024" => Ok(Decimation::DEC_1024),
            "8192" => Ok(Decimation::DEC_8192),
            "65536" => Ok(Decimation::DEC_65536),
            decimation => Err(format!("Unknow decimation '{}'", decimation)),
        }
    }
}

impl ::std::convert::Into<SamplingRate> for Decimation {
    fn into(self) -> SamplingRate {
        match self {
            Decimation::DEC_1 => SamplingRate::RATE_125MHz,
            Decimation::DEC_8 => SamplingRate::RATE_15_6MHz,
            Decimation::DEC_64 => SamplingRate::RATE_1_9MHz,
            Decimation::DEC_1024 => SamplingRate::RATE_103_8kHz,
            Decimation::DEC_8192 => SamplingRate::RATE_15_2kHz,
            Decimation::DEC_65536 => SamplingRate::RATE_1_9kHz,
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum SamplingRate {
    RATE_125MHz,
    RATE_15_6MHz,
    RATE_1_9MHz,
    RATE_103_8kHz,
    RATE_15_2kHz,
    RATE_1_9kHz,
}

impl SamplingRate {
    pub fn get_buffer_duration(self) -> ::std::time::Duration {
        let (s, ns) = match self {
            SamplingRate::RATE_125MHz => (0, 131_072),
            SamplingRate::RATE_15_6MHz => (0, 1_049_000),
            SamplingRate::RATE_1_9MHz => (0, 8_389_000),
            SamplingRate::RATE_103_8kHz => (0, 134_218_000),
            SamplingRate::RATE_15_2kHz => (1, 740_000_000),
            SamplingRate::RATE_1_9kHz => (8, 590_000_000),
        };

        ::std::time::Duration::new(s, ns)
    }
}

impl ::std::fmt::Display for SamplingRate {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let display = match self {
            &SamplingRate::RATE_125MHz => "125 MHz",
            &SamplingRate::RATE_15_6MHz => "15.6 MHz",
            &SamplingRate::RATE_1_9MHz => "1.9 MHz",
            &SamplingRate::RATE_103_8kHz => "103.8 kHz",
            &SamplingRate::RATE_15_2kHz => "15.2 kHz",
            &SamplingRate::RATE_1_9kHz => "1.9 kHz",
        };

        write!(f, "{}", display)
    }
}

impl ::std::convert::Into<String> for SamplingRate {
    fn into(self) -> String {
        let s = match self {
            SamplingRate::RATE_125MHz => "125MHz",
            SamplingRate::RATE_15_6MHz => "15_6MHz",
            SamplingRate::RATE_1_9MHz => "1_9MHz",
            SamplingRate::RATE_103_8kHz => "103_8kHz",
            SamplingRate::RATE_15_2kHz => "15_2kHz",
            SamplingRate::RATE_1_9kHz => "1_9kHz",
        };

        String::from(s)
    }
}

impl ::std::str::FromStr for SamplingRate {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "125000000 Hz" => Ok(SamplingRate::RATE_125MHz),
            "15600000 Hz" => Ok(SamplingRate::RATE_15_6MHz),
            "1900000 Hz" => Ok(SamplingRate::RATE_1_9MHz),
            "103800 Hz" => Ok(SamplingRate::RATE_103_8kHz),
            "15200 Hz" => Ok(SamplingRate::RATE_15_2kHz),
            "1900 Hz" => Ok(SamplingRate::RATE_1_9kHz),
            rate => Err(format!("Unknow sampling rate {}", rate)),
        }
    }
}

#[derive(Clone)]
pub struct Acquire {
    socket: ::std::cell::RefCell<Socket>,
    started: bool,
}

impl ::Module for Acquire {
    fn get_socket<'a>(&'a self) -> ::std::cell::RefMut<'a, ::socket::Socket> {
        self.socket.borrow_mut()
    }
}

impl Acquire {
    pub fn new(socket: Socket) -> Self {
        Acquire {
            socket: ::std::cell::RefCell::new(socket),
            started: false,
        }
    }

    /**
     * Starts acquisition.
     */
    pub fn start(&mut self) {
        self.send("ACQ:START");
        self.started = true;
    }

    /**
     * Stops acquisition.
     */
    pub fn stop(&mut self) {
        self.send("ACQ:STOP");
        self.started = false;
    }

    pub fn is_started(&self) -> bool {
        self.started
    }

    /**
     * Stops acquisition and sets all parameters to default values.
     */
    pub fn reset(&self) {
        self.send("ACQ:RST");
    }

    /**
     * Set decimation factor.
     */
    pub fn set_decimation(&self, decimation: Decimation) {
        self.send(format!("ACQ:DEC {}", Into::<String>::into(decimation)));
    }

    /**
     * Get decimation factor.
     */
    pub fn get_decimation(&self) -> Result<Decimation, String> {
        self.send("ACQ:DEC?");

        self.receive()
            .parse()
    }

    /**
     * Get sampling rate.
     *
     * @warning Calling this command makes buffer overflow.
     * See https://github.com/RedPitaya/RedPitaya/pull/110
     */
    pub fn get_sampling_rate(&self) -> Result<SamplingRate, String> {
        self.send("ACQ:SRAT?");

        self.receive()
            .parse()
    }

    /**
     * Enable averaging.
     */
    pub fn enable_average(&self) {
        self.send("ACQ:AVG ON");
    }

    /**
     * Disable averaging.
     */
    pub fn disable_average(&self) {
        self.send("ACQ:AVG OFF");
    }

    /**
     * Get averaging status.
     */
    pub fn is_average_enabled(&self) -> bool {
        self.send("ACQ:AVG?");

        let message = self.receive();

        match message.as_str() {
            "ON" => true,
            _ => false,
        }
    }

    /**
     * Set gain settings to HIGH or LOW.
     *
     * This gain is referring to jumper settings on Red Pitaya fast analog inputs.
     */
    pub fn set_gain(&self, source: Source, gain: Gain) {
        self.send(format!("ACQ:{}:GAIN {}", Into::<String>::into(source), Into::<String>::into(gain)));
    }

    /**
     * Get gain settings to HIGH or LOW.
     */
    pub fn get_gain(&self, source: Source) -> Result<Gain, String> {
        self.send(format!("ACQ:{}:GAIN?", Into::<String>::into(source)));

        self.receive()
            .parse()
    }
}

#[cfg(test)]
mod test {
    macro_rules! acquire_assert {
        ($f:ident, $e:expr) => {
            let (rx, acquire) = create_acquire();

            acquire.$f();
            assert_eq!($e, rx.recv().unwrap());
        }
    }

    #[test]
    fn test_sampling_rate_get_buffer_duration() {
        let duration = ::std::time::Duration::new(8, 590_000_000);

        assert_eq!(duration, ::acquire::SamplingRate::RATE_1_9kHz.get_buffer_duration());
    }

    #[test]
    fn test_start() {
        let (rx, mut acquire) = create_acquire();

        acquire.start();
        assert_eq!("ACQ:START\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_stop() {
        let (rx, mut acquire) = create_acquire();

        acquire.stop();
        assert_eq!("ACQ:STOP\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_is_started() {
        let (_, mut acquire) = create_acquire();

        assert_eq!(acquire.is_started(), false);
        acquire.start();
        assert_eq!(acquire.is_started(), true);
        acquire.stop();
        assert_eq!(acquire.is_started(), false);
    }

    #[test]
    fn test_reset() {
        acquire_assert!(reset, "ACQ:RST\r\n");
    }

    #[test]
    fn test_set_decimation() {
        let (rx, acquire) = create_acquire();

        acquire.set_decimation(::acquire::Decimation::DEC_8);
        assert_eq!("ACQ:DEC 8\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_decimation() {
        let (_, acquire) = create_acquire();

        assert_eq!(acquire.get_decimation(), Ok(::acquire::Decimation::DEC_1));
    }

    #[test]
    fn test_get_sampling_rate() {
        let (_, acquire) = create_acquire();

        assert_eq!(acquire.get_sampling_rate(), Ok(::acquire::SamplingRate::RATE_125MHz));
    }

    #[test]
    fn test_enable_average() {
        acquire_assert!(enable_average, "ACQ:AVG ON\r\n");
    }

    #[test]
    fn test_disable_average() {
        acquire_assert!(disable_average, "ACQ:AVG OFF\r\n");
    }

    #[test]
    fn test_is_average_enabled() {
        let (_, acquire) = create_acquire();

        assert_eq!(acquire.is_average_enabled(), true);
    }

    #[test]
    fn test_set_gain() {
        let (rx, acquire) = create_acquire();

        acquire.set_gain(::acquire::Source::IN1, ::acquire::Gain::LV);
        assert_eq!("ACQ:SOUR1:GAIN LV\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_gain() {
        let (_, acquire) = create_acquire();

        assert_eq!(acquire.get_gain(::acquire::Source::IN1), Ok(::acquire::Gain::HV));
    }

    fn create_acquire() -> (::std::sync::mpsc::Receiver<String>, ::acquire::Acquire) {
        let (addr, rx) = ::test::launch_server();
        let socket = ::socket::Socket::new(addr);

        (rx, ::acquire::Acquire::new(socket))
    }
}
<|endoftext|>


<commit_before>use std::str::FromStr;
use std::time::Duration;

use clap::{App, AppSettings, Arg};

use crate::app::{SUPPORTED_OUTPUT_FORMATS, SUPPORTED_RECORD_TYPES};
use crate::estimate::Estimate;
use crate::nameserver::predefined;
use crate::resolver::{Lookups, MultiQuery, ResolverGroup, ResolverGroupOpts, ResolverOpts};
use crate::statistics::Statistics;
use std::collections::HashMap;

pub fn setup_clap() -> App<'static, 'static> {
    App::new("mhost")
        .version(env!("CARGO_PKG_VERSION"))
        .author(env!("CARGO_PKG_AUTHORS"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .setting(AppSettings::DisableHelpSubcommand)
        .setting(AppSettings::DeriveDisplayOrder)
        .setting(AppSettings::UnifiedHelpMessage)
        .setting(AppSettings::VersionlessSubcommands)
        .arg(Arg::with_name("domain name")
            .required_unless("list-predefined")
            .index(1)
            .value_name("NAME | IP ADDR | CIDR BLOCK")
            .next_line_help(false)
            .help("domain name, IP address, or CIDR block to lookup")
            .long_help(
                "* DOMAIN NAME may be any valid DNS name, e.g., lukas.pustina.de
* IP ADDR may be any valid IPv4 or IPv4 address, e.g., 192.168.0.1
* CIDR BLOCK may be any valid IPv4 or IPv6 subnet in CIDR notation, e.g., 192.168.0.1/24
  all valid IP addresses of a CIDR block will be queried for a reverse lookup")
        )
        .arg(Arg::with_name("no-system-resolv-opt")
            .long("no-system-resolv-opt")
            .help("Ignores options set in /etc/resolv.conf")
        )
        .arg(Arg::with_name("no-system-nameservers")
            .long("no-system-nameservers")
            .requires("system nameservers")
            .help("Ignores nameservers from /etc/resolv.conf")
        )
        .arg(Arg::with_name("resolv-conf")
            .long("resolv-conf")
            .value_name("FILE")
            .takes_value(true)
            .help("Uses alternative resolv.conf file")
        )
        .arg(Arg::with_name("system nameservers")
            .short("S")
            .long("system-nameserver")
            .value_name("IP ADDR")
            .takes_value(true)
            .multiple(true)
            .help("Adds system nameserver for system lookups; only IP addresses allowed")
        )
        .arg(Arg::with_name("nameservers")
            .short("s")
            .long("nameserver")
            .value_name("HOSTNAME | IP ADDR")
            .takes_value(true)
            .multiple(true)
            .use_delimiter(true)
            .require_delimiter(true)
            .value_delimiter(";")
            .help("Adds nameserver for lookups")
        )
        .arg(Arg::with_name("predefined")
            .short("p")
            .long("predefined")
            .help("Adds predefined nameservers for lookups")
        )
        .arg(Arg::with_name("predefined-filter")
            .long("predefined-filter")
            .value_name("PROTOCOL")
            .multiple(true)
            .use_delimiter(true)
            .require_delimiter(true)
            .default_value("udp,tcp,https,tls")
            .possible_values(&["udp", "tcp", "https", "tls"])
            .default_value_if("predefined", None, "udp,tcp,https,tls")
            .help("Filters predefined nameservers by protocol")
        )
        .arg(Arg::with_name("list-predefined")
            .long("list-predefined")
            .conflicts_with("domain name")
            .help("Lists all predefined nameservers")
        )
        .arg(Arg::with_name("nameservers-from-file")
            .short("f")
            .long("nameservers-from-file")
            .value_name("FILE")
            .takes_value(true)
            .help("Adds nameserver for lookups from file")
        )
        .arg(Arg::with_name("limit")
            .long("limit")
            .value_name("NUMBER")
            .default_value("100")
            .validator(|str| usize::from_str(&str).map(|_| ()).map_err(|_| "invalid number".to_string()))
            .help("Sets max. numnber of nameservers to query")
        )
        .arg(Arg::with_name("record types")
            .short("t")
            .long("record-type")
            .value_name("RECORD TYPE")
            .takes_value(true)
            .multiple(true)
            .use_delimiter(true)
            .require_delimiter(true)
            .default_value("A,AAAA,MX")
            .possible_values(SUPPORTED_RECORD_TYPES)
            .help("Sets record type to lookup, will be ignored in case of IP address lookup")
        )
        .arg(Arg::with_name("all-record-types")
            .long("all")
            .alias("xmas")
            .help("Enables lookups for all record types")
        )
        .arg(Arg::with_name("randomized-lookup")
            .short("R")
            .long("randomized-lookup")
            .help("Switches into randomize lookup mode: every query will be send just once to a one randomly chosen nameserver. This can be used to distribute queries among the available nameservers.")
        )
        .arg(Arg::with_name("max-concurrent-servers")
            .long("max-concurrent-servers")
            .value_name("NUMBER")
            .default_value("10")
            .validator(|str| usize::from_str(&str).map(|_| ()).map_err(|_| "invalid number".to_string()))
            .help("Sets max. concurrent nameservers")
        )
        .arg(Arg::with_name("max-concurrent-requests")
            .long("max-concurrent-requests")
            .value_name("NUMBER")
            .default_value("5")
            .validator(|str| usize::from_str(&str).map(|_| ()).map_err(|_| "invalid number".to_string()))
            .help("Sets max. concurrent requests per nameserver")
        )
        .arg(Arg::with_name("retries")
            .long("retries")
            .value_name("RETRIES")
            .default_value("0")
            .validator(|str| usize::from_str(&str).map(|_| ()).map_err(|_| "invalid number".to_string()))
            .help("Sets number of retries if first lookup to nameserver fails")
        )
        .arg(Arg::with_name("timeout")
            .long("timeout")
            .value_name("TIMEOUT")
            .default_value("5")
            .validator(|str| u64::from_str(&str).map(|_| ()).map_err(|_| "invalid number".to_string()))
            .help("Sets timeout in seconds for responses")
        )
        .arg(Arg::with_name("wait-multiple-responses")
            .long("wait-multiple-responses")
            .help("Waits until timeout for additional responses from nameservers")
        )
        .arg(Arg::with_name("no-abort-on-error")
            .long("no-abort-on-error")
            .help("Sets do-not-ignore errors from nameservers")
        )
        .arg(Arg::with_name("no-abort-on-timeout")
            .long("no-abort-on-timeout")
            .help("Sets do-not-ignore timeouts from nameservers")
        )
        .arg(Arg::with_name("no-aborts")
            .long("no-aborts")
            .help("Sets do-not-ignore errors and timeouts from nameservers")
        )
        .arg(Arg::with_name("output")
            .short("o")
            .long("output")
            .value_name("FORMAT")
            .takes_value(true)
            .default_value("summary")
            .possible_values(SUPPORTED_OUTPUT_FORMATS)
            .help("Sets the output format for result presentation")
        )
        .arg(Arg::with_name("output-options")
            .long("output-options")
            .value_name("OPTIONS")
            .multiple(true)
            .use_delimiter(true)
            .require_delimiter(true)
            .default_value_if("output", Some("json"), "pretty")
            .default_value_if("output", Some("summary"), "human")
            .help("Sets output options")
            .long_help("* Json: 'pretty': Prettifies output
* Summary: 'human': Uses human readable formatting, 'condensed': Simplifies output")
        )
        .arg(Arg::with_name("show-errors")
            .long("show-errors")
            .conflicts_with("quiet")
            .help("Shows error counts")
        )
        .arg(Arg::with_name("quiet")
            .short("q")
            .long("quiet")
            .help("Does not print anything but results")
        )
        .arg(Arg::with_name("no-color")
            .long("no-color")
            .help("Disables colorful output")
        )
        .arg(Arg::with_name("v")
            .short("v")
            .multiple(true)
            .help("Sets the level of verbosity"))
}

pub fn list_predefined_nameservers() {
    println!("List of predefined servers:");
    for ns in predefined::nameserver_configs() {
        println!("* {}", ns);
    }
}

pub fn print_opts(group_opts: &ResolverGroupOpts, opts: &ResolverOpts) {
    println!(
        "Nameservers options: concurrent nameservers={}, max. nameservers={}, concurrent requests={}, retries={}, timeout={} s{}{}{}",
        group_opts.max_concurrent,
        group_opts.limit.unwrap(), // Safe unwrap, because of Clap's default value
        opts.max_concurrent_requests,
        opts.retries,
        opts.timeout.as_secs(),
        if opts.expects_multiple_responses {
            ", wait for additional responses"
        } else {
            ""
        },
        if opts.abort_on_error { ", abort on error" } else { "" },
        if opts.abort_on_timeout {
            ", abort on timeout"
        } else {
            ""
        },
    )
}

pub fn print_estimates(resolvers: &ResolverGroup, query: &MultiQuery) {
    let num_servers = resolvers.opts.limit.unwrap().min(resolvers.len()); // Safe unwrap, because of Clap's default value
    let num_names = query.num_names();
    let num_record_types = query.num_record_types();
    let estimate = resolvers.estimate(query);

    let queries_str = if estimate.min_requests == estimate.max_requests {
        format!(
            "{} {}",
            estimate.min_requests,
            if estimate.min_requests > 1 {
                "requests"
            } else {
                "request"
            }
        )
    } else {
        format!(
            "between {} and {} requests",
            estimate.min_requests, estimate.max_requests
        )
    };
    let namesservers_str = if num_servers > 1 {
        format!("{} nameservers", num_servers)
    } else {
        "1 nameserver".to_string()
    };
    let record_types_str = if num_record_types > 1 {
        format!("{} record types", num_record_types)
    } else {
        "1 record type".to_string()
    };
    let names_str = if num_names > 1 {
        format!("{} record types", num_names)
    } else {
        "1 name".to_string()
    };

    println!(
        "Sending {} to {} for {} of {}.",
        queries_str, namesservers_str, record_types_str, names_str
    )
}

pub fn print_statistics(lookups: &Lookups, total_run_time: Duration) {
    let statistics = lookups.statistics();
    println!(
        "Received {} within {} ms of total run time.",
        statistics,
        total_run_time.as_millis()
    );
}

pub fn print_error_counts(lookups: &Lookups) {
    let mut counts: HashMap<String, usize> = HashMap::new();

    for err in lookups.iter().map(|l| l.result().err()).flatten() {
        let key = format!("{:?}", err);
        let val = counts.entry(key).or_insert(0);
        *val += 1;
    }

    println!("Error counts");
    if counts.is_empty() {
        println!("* No errors occurred.")
    } else {
        for (k, v) in counts.iter() {
            println!("* Err {} occurred {} times", k, v);
        }
    }
}
<commit_msg>Minor fixes<commit_after>use std::str::FromStr;
use std::time::Duration;

use clap::{App, AppSettings, Arg};

use crate::app::{SUPPORTED_OUTPUT_FORMATS, SUPPORTED_RECORD_TYPES};
use crate::estimate::Estimate;
use crate::nameserver::predefined;
use crate::resolver::{Lookups, MultiQuery, ResolverGroup, ResolverGroupOpts, ResolverOpts};
use crate::statistics::Statistics;
use std::collections::HashMap;

pub fn setup_clap() -> App<'static, 'static> {
    App::new("mhost")
        .version(env!("CARGO_PKG_VERSION"))
        .author(env!("CARGO_PKG_AUTHORS"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .setting(AppSettings::DisableHelpSubcommand)
        .setting(AppSettings::DeriveDisplayOrder)
        .setting(AppSettings::UnifiedHelpMessage)
        .setting(AppSettings::VersionlessSubcommands)
        .arg(Arg::with_name("domain name")
            .required_unless("list-predefined")
            .index(1)
            .value_name("NAME | IP ADDR | CIDR BLOCK")
            .next_line_help(false)
            .help("domain name, IP address, or CIDR block to lookup")
            .long_help(
                "* DOMAIN NAME may be any valid DNS name, e.g., lukas.pustina.de
* IP ADDR may be any valid IPv4 or IPv4 address, e.g., 192.168.0.1
* CIDR BLOCK may be any valid IPv4 or IPv6 subnet in CIDR notation, e.g., 192.168.0.1/24
  all valid IP addresses of a CIDR block will be queried for a reverse lookup")
        )
        .arg(Arg::with_name("no-system-resolv-opt")
            .long("no-system-resolv-opt")
            .help("Ignores options set in /etc/resolv.conf")
        )
        .arg(Arg::with_name("no-system-nameservers")
            .long("no-system-nameservers")
            .requires("system nameservers")
            .help("Ignores nameservers from /etc/resolv.conf")
        )
        .arg(Arg::with_name("resolv-conf")
            .long("resolv-conf")
            .value_name("FILE")
            .takes_value(true)
            .help("Uses alternative resolv.conf file")
        )
        .arg(Arg::with_name("system nameservers")
            .short("S")
            .long("system-nameserver")
            .value_name("IP ADDR")
            .takes_value(true)
            .multiple(true)
            .help("Adds system nameserver for system lookups; only IP addresses allowed")
        )
        .arg(Arg::with_name("nameservers")
            .short("s")
            .long("nameserver")
            .value_name("HOSTNAME | IP ADDR")
            .takes_value(true)
            .multiple(true)
            .use_delimiter(true)
            .require_delimiter(true)
            .value_delimiter(";")
            .help("Adds nameserver for lookups")
        )
        .arg(Arg::with_name("predefined")
            .short("p")
            .long("predefined")
            .help("Adds predefined nameservers for lookups")
        )
        .arg(Arg::with_name("predefined-filter")
            .long("predefined-filter")
            .value_name("PROTOCOL")
            .multiple(true)
            .use_delimiter(true)
            .require_delimiter(true)
            .default_value("udp,tcp,https,tls")
            .possible_values(&["udp", "tcp", "https", "tls"])
            .default_value_if("predefined", None, "udp,tcp,https,tls")
            .requires("predefined")
            .help("Filters predefined nameservers by protocol")
        )
        .arg(Arg::with_name("list-predefined")
            .long("list-predefined")
            .conflicts_with("domain name")
            .help("Lists all predefined nameservers")
        )
        .arg(Arg::with_name("nameservers-from-file")
            .short("f")
            .long("nameservers-from-file")
            .value_name("FILE")
            .takes_value(true)
            .help("Adds nameserver for lookups from file")
        )
        .arg(Arg::with_name("limit")
            .long("limit")
            .value_name("NUMBER")
            .default_value("100")
            .validator(|str| usize::from_str(&str).map(|_| ()).map_err(|_| "invalid number".to_string()))
            .help("Sets max. number of nameservers to query")
        )
        .arg(Arg::with_name("record types")
            .short("t")
            .long("record-type")
            .value_name("RECORD TYPE")
            .takes_value(true)
            .multiple(true)
            .use_delimiter(true)
            .require_delimiter(true)
            .default_value("A,AAAA,MX")
            .possible_values(SUPPORTED_RECORD_TYPES)
            .help("Sets record type to lookup, will be ignored in case of IP address lookup")
        )
        .arg(Arg::with_name("all-record-types")
            .long("all")
            .alias("xmas")
            .help("Enables lookups for all record types")
        )
        .arg(Arg::with_name("randomized-lookup")
            .short("R")
            .long("randomized-lookup")
            .help("Switches into randomize lookup mode: every query will be send just once to a one randomly chosen nameserver. This can be used to distribute queries among the available nameservers.")
        )
        .arg(Arg::with_name("max-concurrent-servers")
            .long("max-concurrent-servers")
            .value_name("NUMBER")
            .default_value("10")
            .validator(|str| usize::from_str(&str).map(|_| ()).map_err(|_| "invalid number".to_string()))
            .help("Sets max. concurrent nameservers")
        )
        .arg(Arg::with_name("max-concurrent-requests")
            .long("max-concurrent-requests")
            .value_name("NUMBER")
            .default_value("5")
            .validator(|str| usize::from_str(&str).map(|_| ()).map_err(|_| "invalid number".to_string()))
            .help("Sets max. concurrent requests per nameserver")
        )
        .arg(Arg::with_name("retries")
            .long("retries")
            .value_name("RETRIES")
            .default_value("0")
            .validator(|str| usize::from_str(&str).map(|_| ()).map_err(|_| "invalid number".to_string()))
            .help("Sets number of retries if first lookup to nameserver fails")
        )
        .arg(Arg::with_name("timeout")
            .long("timeout")
            .value_name("TIMEOUT")
            .default_value("5")
            .validator(|str| u64::from_str(&str).map(|_| ()).map_err(|_| "invalid number".to_string()))
            .help("Sets timeout in seconds for responses")
        )
        .arg(Arg::with_name("wait-multiple-responses")
            .long("wait-multiple-responses")
            .help("Waits until timeout for additional responses from nameservers")
        )
        .arg(Arg::with_name("no-abort-on-error")
            .long("no-abort-on-error")
            .help("Sets do-not-ignore errors from nameservers")
        )
        .arg(Arg::with_name("no-abort-on-timeout")
            .long("no-abort-on-timeout")
            .help("Sets do-not-ignore timeouts from nameservers")
        )
        .arg(Arg::with_name("no-aborts")
            .long("no-aborts")
            .help("Sets do-not-ignore errors and timeouts from nameservers")
        )
        .arg(Arg::with_name("output")
            .short("o")
            .long("output")
            .value_name("FORMAT")
            .takes_value(true)
            .default_value("summary")
            .possible_values(SUPPORTED_OUTPUT_FORMATS)
            .help("Sets the output format for result presentation")
        )
        .arg(Arg::with_name("output-options")
            .long("output-options")
            .value_name("OPTIONS")
            .multiple(true)
            .use_delimiter(true)
            .require_delimiter(true)
            .default_value_if("output", Some("json"), "pretty")
            .default_value_if("output", Some("summary"), "human")
            .help("Sets output options")
            .long_help("* Json: 'pretty': Prettifies output
* Summary: 'human': Uses human readable formatting, 'condensed': Simplifies output")
        )
        .arg(Arg::with_name("show-errors")
            .long("show-errors")
            .conflicts_with("quiet")
            .help("Shows error counts")
        )
        .arg(Arg::with_name("quiet")
            .short("q")
            .long("quiet")
            .help("Does not print anything but results")
        )
        .arg(Arg::with_name("no-color")
            .long("no-color")
            .help("Disables colorful output")
        )
        .arg(Arg::with_name("v")
            .short("v")
            .multiple(true)
            .help("Sets the level of verbosity"))
}

pub fn list_predefined_nameservers() {
    println!("List of predefined servers:");
    for ns in predefined::nameserver_configs() {
        println!("* {}", ns);
    }
}

pub fn print_opts(group_opts: &ResolverGroupOpts, opts: &ResolverOpts) {
    println!(
        "Nameservers options: concurrent nameservers={}, max. nameservers={}, concurrent requests={}, retries={}, timeout={} s{}{}{}",
        group_opts.max_concurrent,
        group_opts.limit.unwrap(), // Safe unwrap, because of Clap's default value
        opts.max_concurrent_requests,
        opts.retries,
        opts.timeout.as_secs(),
        if opts.expects_multiple_responses {
            ", wait for additional responses"
        } else {
            ""
        },
        if opts.abort_on_error { ", abort on error" } else { "" },
        if opts.abort_on_timeout {
            ", abort on timeout"
        } else {
            ""
        },
    )
}

pub fn print_estimates(resolvers: &ResolverGroup, query: &MultiQuery) {
    let num_servers = resolvers.opts.limit.unwrap().min(resolvers.len()); // Safe unwrap, because of Clap's default value
    let num_names = query.num_names();
    let num_record_types = query.num_record_types();
    let estimate = resolvers.estimate(query);

    let queries_str = if estimate.min_requests == estimate.max_requests {
        format!(
            "{} {}",
            estimate.min_requests,
            if estimate.min_requests > 1 {
                "requests"
            } else {
                "request"
            }
        )
    } else {
        format!(
            "between {} and {} requests",
            estimate.min_requests, estimate.max_requests
        )
    };
    let namesservers_str = if num_servers > 1 {
        format!("{} nameservers", num_servers)
    } else {
        "1 nameserver".to_string()
    };
    let record_types_str = if num_record_types > 1 {
        format!("{} record types", num_record_types)
    } else {
        "1 record type".to_string()
    };
    let names_str = if num_names > 1 {
        format!("{} record types", num_names)
    } else {
        "1 name".to_string()
    };

    println!(
        "Sending {} to {} for {} of {}.",
        queries_str, namesservers_str, record_types_str, names_str
    )
}

pub fn print_statistics(lookups: &Lookups, total_run_time: Duration) {
    let statistics = lookups.statistics();
    println!(
        "Received {} within {} ms of total run time.",
        statistics,
        total_run_time.as_millis()
    );
}

pub fn print_error_counts(lookups: &Lookups) {
    let mut counts: HashMap<String, usize> = HashMap::new();

    for err in lookups.iter().map(|l| l.result().err()).flatten() {
        let key = format!("{:?}", err);
        let val = counts.entry(key).or_insert(0);
        *val += 1;
    }

    println!("Error counts");
    if counts.is_empty() {
        println!("* No errors occurred.")
    } else {
        for (k, v) in counts.iter() {
            println!("* Err {} occurred {} times", k, v);
        }
    }
}
<|endoftext|>