From c6e38690beb44eabd96cfbbd74de475982f09270 Mon Sep 17 00:00:00 2001
From: Lucas Bajolet <lucas.bajolet@hashicorp.com>
Date: Mon, 6 May 2024 16:44:09 -0400
Subject: [PATCH] test: add LineCount/Compare pipe/tester

As a common use case in console-oriented pipelines, we check that a
specific command returned a certain number of lines.
With the combination of LineCount and Compare, we can do exactly this.
---
 test/pipe_checker_test.go | 79 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)

diff --git a/test/pipe_checker_test.go b/test/pipe_checker_test.go
index 8b2f6331e..72193a9e9 100644
--- a/test/pipe_checker_test.go
+++ b/test/pipe_checker_test.go
@@ -3,6 +3,7 @@ package test
 import (
 	"fmt"
 	"regexp"
+	"strconv"
 	"strings"
 )
 
@@ -31,6 +32,18 @@ func PipeGrep(expression string) Pipe {
 	})
 }
 
+// LineCount counts the number of lines received
+//
+// This excludes empty lines.
+func LineCount() Pipe {
+	return CustomPipe(func(s string) (string, error) {
+		lines := strings.FieldsFunc(s, func(r rune) bool {
+			return r == '\n'
+		})
+		return fmt.Sprintf("%d\n", len(lines)), nil
+	})
+}
+
 // Tester is the end of a pipe for testing purposes.
 //
 // Once multiple commands have been piped together in a pipeline, we can
@@ -69,6 +82,72 @@ func EmptyInput() Tester {
 	})
 }
 
+type op int
+
+const (
+	eq op = iota
+	ne
+	gt
+	ge
+	lt
+	le
+)
+
+func (op op) String() string {
+	switch op {
+	case eq:
+		return "=="
+	case ne:
+		return "!="
+	case gt:
+		return ">"
+	case ge:
+		return ">="
+	case lt:
+		return "<"
+	case le:
+		return "<="
+	}
+
+	panic(fmt.Sprintf("Unknown operator %d", op))
+}
+
+// IntCompare reads the input from the pipeline and compares it to a value using `op`
+//
+// If the input is not an int, this fails.
+func IntCompare(op op, value int) Tester {
+	return CustomTester(func(in string) error {
+		n, err := strconv.Atoi(in)
+		if err != nil {
+			return fmt.Errorf("not an integer %q: %s", in, err)
+		}
+
+		var result bool
+		switch op {
+		case eq:
+			result = n == value
+		case ne:
+			result = n != value
+		case gt:
+			result = n > value
+		case ge:
+			result = n >= value
+		case lt:
+			result = n < value
+		case le:
+			result = n <= value
+		default:
+			panic(fmt.Sprintf("Unsupported operator %d, make sure the operation is implemented for IntCompare", op))
+		}
+
+		if !result {
+			return fmt.Errorf("comparison failed: %d %s %d -> %t", n, op, value, result)
+		}
+
+		return nil
+	})
+}
+
 // PipeChecker is a kind of checker that essentially lets users write mini
 // gadgets that pipe inputs together, and compose those to end as a true/false
 // statement, which translates to an error.