xan 0.57.1

use std::borrow::ToOwned;

use crate::workdir::Workdir;

macro_rules! part_eq {
    ($wrk:expr, $path:expr, $expected:expr) => {
        assert_eq!(
            $wrk.from_str::<String>(&$wrk.path($path)),
            $expected.to_owned()
        );
    };
}

fn data(headers: bool) -> Vec<Vec<String>> {
    let mut rows = vec![
        svec!["NY", "Manhattan"],
        svec!["CA", "San Francisco"],
        svec!["TX", "Dallas"],
        svec!["NY", "Buffalo"],
        svec!["TX", "Fort Worth"],
    ];
    if headers {
        rows.insert(0, svec!["state", "city"]);
    }
    rows
}

#[test]
fn partition() {
    let wrk = Workdir::new("partition");
    wrk.create("in.csv", data(true));

    let mut cmd = wrk.command("partition");
    cmd.arg("state")
        .arg("-C")
        .arg("-O")
        .arg(&wrk.path("."))
        .arg("in.csv");
    wrk.run(&mut cmd);

    part_eq!(
        wrk,
        "CA.csv",
        "\
state,city
CA,San Francisco
"
    );
    part_eq!(
        wrk,
        "NY.csv",
        "\
state,city
NY,Manhattan
NY,Buffalo
"
    );
    part_eq!(
        wrk,
        "TX.csv",
        "\
state,city
TX,Dallas
TX,Fort Worth
"
    );
}

#[test]
fn partition_case_insensitive() {
    let wrk = Workdir::new("partition_case_insensitive");
    wrk.create(
        "in.csv",
        vec![
            svec!["state", "city"],
            svec!["Co", "Manhattan"],
            svec!["co", "San Francisco"],
            svec!["Co", "Dallas"],
            svec!["CO", "Buffalo"],
            svec!["VT", "Fort Worth"],
        ],
    );

    let mut cmd = wrk.command("partition");
    cmd.arg("state").arg("-O").arg(&wrk.path(".")).arg("in.csv");
    wrk.run(&mut cmd);

    part_eq!(
        wrk,
        "Co.csv",
        "\
state,city
Co,Manhattan
Co,Dallas
"
    );
    part_eq!(
        wrk,
        "co_1.csv",
        "\
state,city
co,San Francisco
"
    );
    part_eq!(
        wrk,
        "CO_2.csv",
        "\
state,city
CO,Buffalo
"
    );
    part_eq!(
        wrk,
        "VT.csv",
        "\
state,city
VT,Fort Worth
"
    );
}

#[test]
fn partition_case_insensitive_sorted() {
    let wrk = Workdir::new("partition_case_insensitive");
    wrk.create(
        "in.csv",
        vec![
            svec!["state", "city"],
            svec!["Co", "Manhattan"],
            svec!["Co", "Dallas"],
            svec!["co", "San Francisco"],
            svec!["CO", "Buffalo"],
            svec!["VT", "Fort Worth"],
        ],
    );

    let mut cmd = wrk.command("partition");
    cmd.arg("state")
        .arg("-O")
        .arg(&wrk.path("."))
        .arg("--sorted")
        .arg("in.csv");
    wrk.run(&mut cmd);

    part_eq!(
        wrk,
        "Co.csv",
        "\
state,city
Co,Manhattan
Co,Dallas
"
    );
    part_eq!(
        wrk,
        "co_1.csv",
        "\
state,city
co,San Francisco
"
    );
    part_eq!(
        wrk,
        "CO_2.csv",
        "\
state,city
CO,Buffalo
"
    );
    part_eq!(
        wrk,
        "VT.csv",
        "\
state,city
VT,Fort Worth
"
    );
}

#[test]
fn partition_drop() {
    let wrk = Workdir::new("partition_drop");
    wrk.create("in.csv", data(true));

    let mut cmd = wrk.command("partition");
    cmd.arg("--drop")
        .arg("state")
        .arg("-C")
        .arg("-O")
        .arg(&wrk.path("."))
        .arg("in.csv");
    wrk.run(&mut cmd);

    part_eq!(
        wrk,
        "CA.csv",
        "\
city
San Francisco
"
    );
    part_eq!(
        wrk,
        "NY.csv",
        "\
city
Manhattan
Buffalo
"
    );
    part_eq!(
        wrk,
        "TX.csv",
        "\
city
Dallas
Fort Worth
"
    );
}

#[test]
fn partition_without_headers() {
    let wrk = Workdir::new("partition_without_headers");
    wrk.create("in.csv", data(false));

    let mut cmd = wrk.command("partition");
    cmd.arg("--no-headers")
        .arg("0")
        .arg("-C")
        .arg("-O")
        .arg(&wrk.path("."))
        .arg("in.csv");
    wrk.run(&mut cmd);

    part_eq!(
        wrk,
        "CA.csv",
        "\
CA,San Francisco
"
    );
    part_eq!(
        wrk,
        "NY.csv",
        "\
NY,Manhattan
NY,Buffalo
"
    );
    part_eq!(
        wrk,
        "TX.csv",
        "\
TX,Dallas
TX,Fort Worth
"
    );
}

#[test]
fn partition_drop_without_headers() {
    let wrk = Workdir::new("partition_drop_without_headers");
    wrk.create("in.csv", data(false));

    let mut cmd = wrk.command("partition");
    cmd.arg("--drop")
        .arg("--no-headers")
        .arg("0")
        .arg("-C")
        .arg("-O")
        .arg(&wrk.path("."))
        .arg("in.csv");
    wrk.run(&mut cmd);

    part_eq!(
        wrk,
        "CA.csv",
        "\
San Francisco
"
    );
    part_eq!(
        wrk,
        "NY.csv",
        "\
Manhattan
Buffalo
"
    );
    part_eq!(
        wrk,
        "TX.csv",
        "\
Dallas
Fort Worth
"
    );
}

#[test]
fn partition_into_new_directory() {
    let wrk = Workdir::new("partition_into_new_directory");
    wrk.create("in.csv", data(true));

    let mut cmd = wrk.command("partition");
    cmd.arg("state")
        .arg("-C")
        .arg("-O")
        .arg(&wrk.path("out"))
        .arg("in.csv");
    wrk.run(&mut cmd);

    assert!(wrk.path("out/NY.csv").exists());
}

#[test]
fn partition_custom_filename() {
    let wrk = Workdir::new("partition_custom_filename");
    wrk.create("in.csv", data(true));

    let mut cmd = wrk.command("partition");
    cmd.args(["--filename", "state-{}-partition.csv"])
        .arg("state")
        .arg("-C")
        .arg("-O")
        .arg(&wrk.path("."))
        .arg("in.csv");
    wrk.run(&mut cmd);

    assert!(wrk.path("state-NY-partition.csv").exists());
}

#[test]
fn partition_custom_filename_with_directory() {
    let wrk = Workdir::new("partition_custom_filename_with_directory");
    wrk.create("in.csv", data(true));

    let mut cmd = wrk.command("partition");
    cmd.args(["--filename", "{}/cities.csv"])
        .arg("state")
        .arg("-C")
        .arg("-O")
        .arg(&wrk.path("."))
        .arg("in.csv");
    wrk.run(&mut cmd);

    // This variation also helps with parallel partition jobs.
    assert!(wrk.path("NY/cities.csv").exists());
}

#[test]
fn partition_invalid_filename() {
    let wrk = Workdir::new("partition_invalid_filename");
    wrk.create("in.csv", data(true));

    let mut cmd = wrk.command("partition");
    cmd.args(["--filename", "foo.csv"])
        .arg("state")
        .arg("-C")
        .arg("-O")
        .arg(&wrk.path("."))
        .arg("in.csv");
    wrk.assert_err(&mut cmd);

    let mut cmd = wrk.command("partition");
    cmd.args(["--filename", "{}{}.csv"])
        .arg("state")
        .arg("-C")
        .arg("-O")
        .arg(&wrk.path("."))
        .arg("in.csv");
    wrk.assert_err(&mut cmd);
}

fn tricky_data() -> Vec<Vec<String>> {
    vec![
        svec!["key", "explanation"],
        svec!["", "empty key"],
        svec!["empty", "the string empty"],
        svec!["unsafe _1$!,\"", "unsafe in shell"],
        svec!["collision", "ordinary value"],
        svec!["collision", "in same file"],
        svec!["coll ision", "collides"],
        svec!["collision!", "collides again"],
        svec!["collision_2", "collides with disambiguated"],
    ]
}

#[test]
fn partition_with_tricky_key_values() {
    let wrk = Workdir::new("partition_with_tricky_key_values");
    wrk.create("in.csv", tricky_data());

    let mut cmd = wrk.command("partition");
    cmd.arg("key")
        .arg("-C")
        .arg("-O")
        .arg(&wrk.path("."))
        .arg("in.csv");
    wrk.run(&mut cmd);

    part_eq!(
        wrk,
        "empty.csv",
        "\
key,explanation
,empty key
"
    );
    part_eq!(
        wrk,
        "empty_1.csv",
        "\
key,explanation
empty,the string empty
"
    );
    part_eq!(
        wrk,
        "unsafe_1.csv",
        r#"key,explanation
"unsafe _1$!,""",unsafe in shell
"#
    );
    part_eq!(
        wrk,
        "collision.csv",
        "\
key,explanation
collision,ordinary value
collision,in same file
"
    );
    part_eq!(
        wrk,
        "collision_2.csv",
        "\
key,explanation
coll ision,collides
"
    );
    part_eq!(
        wrk,
        "collision_3.csv",
        "\
key,explanation
collision!,collides again
"
    );
    // Tricky! We didn't see this an input, but we did generate it as an
    // output already.
    part_eq!(
        wrk,
        "collision_2_4.csv",
        "\
key,explanation
collision_2,collides with disambiguated
"
    );
}

fn prefix_data() -> Vec<Vec<String>> {
    vec![
        svec!["state", "city"],
        svec!["MA", "Boston"],
        svec!["ME", "Portland"],
        svec!["M", "Too short"],
        svec!["CA", "San Francisco"],
        svec!["CO", "Denver"],
    ]
}

#[test]
fn partition_with_prefix_length() {
    let wrk = Workdir::new("partition_with_prefix_length");
    wrk.create("in.csv", prefix_data());

    let mut cmd = wrk.command("partition");
    cmd.args(["--prefix-length", "1"])
        .arg("state")
        .arg("-C")
        .arg("-O")
        .arg(&wrk.path("."))
        .arg("in.csv");
    wrk.run(&mut cmd);

    part_eq!(
        wrk,
        "M.csv",
        "\
state,city
MA,Boston
ME,Portland
M,Too short
"
    );
    part_eq!(
        wrk,
        "C.csv",
        "\
state,city
CA,San Francisco
CO,Denver
"
    );
}

fn sorted_data(headers: bool) -> Vec<Vec<String>> {
    let mut rows = vec![
        svec!["NY", "Manhattan"],
        svec!["NY", "Buffalo"],
        svec!["CA", "San Francisco"],
        svec!["TX", "Dallas"],
        svec!["TX", "Fort Worth"],
    ];
    if headers {
        rows.insert(0, svec!["state", "city"]);
    }
    rows
}

#[test]
fn partition_sorted() {
    let wrk = Workdir::new("partition_sorted");
    wrk.create("in.csv", sorted_data(true));

    let mut cmd = wrk.command("partition");
    cmd.arg("state")
        .arg("-C")
        .arg("-O")
        .arg(&wrk.path("."))
        .arg("-S")
        .arg("in.csv");
    wrk.run(&mut cmd);

    part_eq!(
        wrk,
        "CA.csv",
        "\
state,city
CA,San Francisco
"
    );
    part_eq!(
        wrk,
        "NY.csv",
        "\
state,city
NY,Manhattan
NY,Buffalo
"
    );
    part_eq!(
        wrk,
        "TX.csv",
        "\
state,city
TX,Dallas
TX,Fort Worth
"
    );
}