Query piped nodes for -R/-F with -C; count input only for bare -C

main_test.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync/atomic"
 	"testing"
 
 	"gopkg.in/yaml.v3"
@@ -872,3 +873,74 @@ func TestRun_AllFacts_PrintsSortedByName(t *testing.T) {
 		t.Fatalf("facts not sorted by name: %v", lines)
 	}
 }
+
+// ---- run: count with piped stdin --------------------------------------------
+
+func withStdin(t *testing.T, content string) {
+	t.Helper()
+	path := filepath.Join(t.TempDir(), "stdin")
+	if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
+		t.Fatal(err)
+	}
+	f, err := os.Open(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	old := os.Stdin
+	os.Stdin = f
+	t.Cleanup(func() { os.Stdin = old; _ = f.Close() })
+}
+
+const pipedLookup = "hosta.main.unkin.net puppetdbapi.service.consul\nhostb.main.unkin.net puppetdbapi.service.consul\n"
+
+func TestRun_CountPipedWithRoleQueriesEachNode(t *testing.T) {
+	withStdin(t, pipedLookup)
+	var queried atomic.Int32
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		queried.Add(1)
+		_ = json.NewEncoder(w).Encode([]fact{{Certname: "h", Name: "enc_role", Value: rawJSON("roles::vault")}})
+	}))
+	t.Cleanup(srv.Close)
+	cfg := config{PuppetDBURL: srv.URL, RoleFact: "enc_role"}
+
+	out := strings.TrimSpace(captureStdout(t, func() {
+		if err := run(cfg, "", "", "", true, false, false, false, false, true, false, false, false); err != nil {
+			t.Fatalf("run returned error: %v", err)
+		}
+	}))
+	if out != "2: roles::vault" {
+		t.Fatalf("expected role counts of piped nodes, got: %q", out)
+	}
+	if n := queried.Load(); n != 2 {
+		t.Fatalf("expected one query per piped node, got %d", n)
+	}
+}
+
+func TestRun_CountPipedWithoutQueryCountsInput(t *testing.T) {
+	withStdin(t, pipedLookup+"hostc.main.unkin.net other\n")
+	cfg := config{PuppetDBURL: "http://unused", RoleFact: "enc_role"}
+
+	out := strings.TrimSpace(captureStdout(t, func() {
+		if err := run(cfg, "", "", "", false, false, false, false, false, true, false, false, false); err != nil {
+			t.Fatalf("run returned error: %v", err)
+		}
+	}))
+	if out != "1: other\n2: puppetdbapi.service.consul" {
+		t.Fatalf("expected counts of piped values, got: %q", out)
+	}
+}
+
+func TestRun_CountWithoutStdinRequiresRoleOrFact(t *testing.T) {
+	devnull, err := os.Open(os.DevNull)
+	if err != nil {
+		t.Fatal(err)
+	}
+	old := os.Stdin
+	os.Stdin = devnull
+	t.Cleanup(func() { os.Stdin = old; _ = devnull.Close() })
+
+	err = run(config{PuppetDBURL: "http://unused"}, "", "", "", false, false, false, false, false, true, false, false, false)
+	if err == nil || !strings.Contains(err.Error(), "-R or -F") {
+		t.Fatalf("expected -R/-F requirement error, got: %v", err)
+	}
+}