1#![deny(warnings)]
2
3use std::io::BufRead;
4use std::iter::Iterator;
5
6mod sampling;
7
8fn main() {
9    // Read the user's command line arguments (if any)
10    //
11    //   0 arguments  = get a single random line
12    //   1 argument k = get that number of lines
13    //  >1 arguments  = error
14    //
15    let args: Vec<_> = std::env::args().collect();
16
17    let k = match args.len() {
18        1 => 1,
19        2 => match args[1].parse::<usize>() {
20            Ok(parsed_k) if parsed_k > 0 => parsed_k,
21            _ => {
22                eprintln!("Usage: randline [k]");
23                std::process::exit(1)
24            }
25        },
26        _ => {
27            eprintln!("Usage: randline [k]");
28            std::process::exit(1)
29        }
30    };
31
32    let lines = std::io::stdin().lock().lines().map(|line| match line {
33        Ok(ln) => ln,
34        Err(e) => {
35            eprintln!("Unable to read from stdin: {:?}", e);
36            std::process::exit(1)
37        }
38    });
39
40    let sample = sampling::reservoir_sample(lines, k);
41
42    for line in sample {
43        println!("{}", line);
44    }
45}
46
47#[cfg(test)]
48mod cli_tests {
49    use assert_cmd::Command;
50
51    // Note: for the purposes of the CLI tests, I trust that the reservoir
52    // sampling code works correctly -- that's tested separately.  I'm just
53    // checking the CLI parses the options correctly, so all the input lines
54    // are the same for easy assertions.
55
56    // If you call `randline` without any arguments, it picks a single line.
57    #[test]
58    fn it_selects_a_single_line_if_no_arg() {
59        Command::cargo_bin("randline")
60            .unwrap()
61            .write_stdin("a\na\na\na\na\na\n")
62            .assert()
63            .success()
64            .stdout("a\n")
65            .stderr("");
66    }
67
68    // If you pass an argument `k` and there are more lines than `k`,
69    // it selects a subset of them.
70    #[test]
71    fn it_selects_k_lines_if_more_lines_than_k() {
72        Command::cargo_bin("randline")
73            .unwrap()
74            .arg("2")
75            .write_stdin("a\na\na\na\na\na\n")
76            .assert()
77            .success()
78            .stdout("a\na\n")
79            .stderr("");
80    }
81
82    // If you pass an argument `k` and there are that number of lines,
83    // it selects all of them.
84    #[test]
85    fn it_selects_k_lines_if_equal_lines_to_k() {
86        Command::cargo_bin("randline")
87            .unwrap()
88            .arg("2")
89            .write_stdin("a\na\n")
90            .assert()
91            .success()
92            .stdout("a\na\n")
93            .stderr("");
94    }
95
96    // If you pass an argument `k` and there are less lines than `k`,
97    // it selects all of them.
98    #[test]
99    fn it_selects_k_lines_if_less_lines_than_k() {
100        Command::cargo_bin("randline")
101            .unwrap()
102            .arg("5")
103            .write_stdin("a\na\n")
104            .assert()
105            .success()
106            .stdout("a\na\n")
107            .stderr("");
108    }
109
110    // Passing a non-integer argument is an error.
111    #[test]
112    fn it_fails_if_non_integer_argument() {
113        Command::cargo_bin("randline")
114            .unwrap()
115            .arg("XXX")
116            .assert()
117            .failure()
118            .code(1)
119            .stdout("")
120            .stderr("Usage: randline [k]\n");
121    }
122
123    // Passing k=0 is an error.
124    #[test]
125    fn it_fails_if_k_equals_zero() {
126        Command::cargo_bin("randline")
127            .unwrap()
128            .arg("0")
129            .assert()
130            .failure()
131            .code(1)
132            .stdout("")
133            .stderr("Usage: randline [k]\n");
134    }
135
136    // Passing k<0 is an error.
137    #[test]
138    fn it_fails_if_k_negative() {
139        Command::cargo_bin("randline")
140            .unwrap()
141            .arg("-1")
142            .assert()
143            .failure()
144            .code(1)
145            .stdout("")
146            .stderr("Usage: randline [k]\n");
147    }
148
149    // Passing more than one argument is an error.
150    #[test]
151    fn it_fails_if_too_many_args() {
152        Command::cargo_bin("randline")
153            .unwrap()
154            .args(&["1", "2", "3"])
155            .assert()
156            .failure()
157            .code(1)
158            .stdout("")
159            .stderr("Usage: randline [k]\n");
160    }
161}