From 69f23c25214f8e42c9556c868fdeed0361a8c00c Mon Sep 17 00:00:00 2001
From: Ben Wiederhake <BenWiederhake.GitHub@gmx.de>
Date: Thu, 15 Feb 2024 22:03:21 +0100
Subject: [PATCH] shuf: obey all headcount args, not just the last

---
 src/uu/shuf/src/shuf.rs    |  1 +
 tests/by-util/test_shuf.rs | 66 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/src/uu/shuf/src/shuf.rs b/src/uu/shuf/src/shuf.rs
index 4052af49e..7df7f1e44 100644
--- a/src/uu/shuf/src/shuf.rs
+++ b/src/uu/shuf/src/shuf.rs
@@ -145,6 +145,7 @@ pub fn uu_app() -> Command {
                 .short('n')
                 .long(options::HEAD_COUNT)
                 .value_name("COUNT")
+                .action(clap::ArgAction::Append)
                 .help("output at most COUNT lines"),
         )
         .arg(
diff --git a/tests/by-util/test_shuf.rs b/tests/by-util/test_shuf.rs
index eca914f9f..91167b36e 100644
--- a/tests/by-util/test_shuf.rs
+++ b/tests/by-util/test_shuf.rs
@@ -131,6 +131,72 @@ fn test_head_count() {
     );
 }
 
+#[test]
+fn test_head_count_multi_big_then_small() {
+    let repeat_limit = 5;
+    let input_seq = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+    let input = input_seq
+        .iter()
+        .map(ToString::to_string)
+        .collect::<Vec<String>>()
+        .join("\n");
+
+    let result = new_ucmd!()
+        .arg("-n")
+        .arg(&(repeat_limit + 1).to_string())
+        .arg("-n")
+        .arg(&repeat_limit.to_string())
+        .pipe_in(input.as_bytes())
+        .succeeds();
+    result.no_stderr();
+
+    let result_seq: Vec<i32> = result
+        .stdout_str()
+        .split('\n')
+        .filter(|x| !x.is_empty())
+        .map(|x| x.parse().unwrap())
+        .collect();
+    assert_eq!(result_seq.len(), repeat_limit, "Output is not limited");
+    assert!(
+        result_seq.iter().all(|x| input_seq.contains(x)),
+        "Output includes element not from input: {}",
+        result.stdout_str()
+    );
+}
+
+#[test]
+fn test_head_count_multi_small_then_big() {
+    let repeat_limit = 5;
+    let input_seq = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+    let input = input_seq
+        .iter()
+        .map(ToString::to_string)
+        .collect::<Vec<String>>()
+        .join("\n");
+
+    let result = new_ucmd!()
+        .arg("-n")
+        .arg(&repeat_limit.to_string())
+        .arg("-n")
+        .arg(&(repeat_limit + 1).to_string())
+        .pipe_in(input.as_bytes())
+        .succeeds();
+    result.no_stderr();
+
+    let result_seq: Vec<i32> = result
+        .stdout_str()
+        .split('\n')
+        .filter(|x| !x.is_empty())
+        .map(|x| x.parse().unwrap())
+        .collect();
+    assert_eq!(result_seq.len(), repeat_limit, "Output is not limited");
+    assert!(
+        result_seq.iter().all(|x| input_seq.contains(x)),
+        "Output includes element not from input: {}",
+        result.stdout_str()
+    );
+}
+
 #[test]
 fn test_repeat() {
     let repeat_limit = 15000;