From 351ef98df3afc2b90eb7c7fa896bbd11d6ce01b5 Mon Sep 17 00:00:00 2001
From: Travis Bischel <travis.bischel+github@gmail.com>
Date: Sat, 30 Dec 2023 23:52:24 -0700
Subject: [PATCH] kgo tests: support 512 in KGO_TEST_SCRAM

---
 pkg/kgo/helpers_test.go | 21 +++++++++++++++++----
 1 file changed, 17 insertions(+), 4 deletions(-)

diff --git a/pkg/kgo/helpers_test.go b/pkg/kgo/helpers_test.go
index 6d90a436..7cff717a 100644
--- a/pkg/kgo/helpers_test.go
+++ b/pkg/kgo/helpers_test.go
@@ -43,7 +43,7 @@ var (
 	// KGO_TEST_TLS: DSL syntax is ({ca|cert|key}:path),{1,3}
 	testCert *tls.Config
 
-	// KGO_TEST_SCRAM: DSL is user:pass; we require SCRAM-SHA-256
+	// KGO_TEST_SCRAM: DSL is user:pass(:num); we assume 256
 	saslScram sasl.Mechanism
 
 	// We create topics with a different number of partitions to exercise
@@ -109,13 +109,26 @@ func init() {
 	}
 	if saslStr, exists := os.LookupEnv("KGO_TEST_SCRAM"); exists {
 		split := strings.Split(saslStr, ":")
-		if len(split) != 2 {
+		if len(split) != 2 && len(split) != 3 {
 			panic(fmt.Sprintf("invalid scram format %q", saslStr))
 		}
-		saslScram = (scram.Auth{
+		a := scram.Auth{
 			User: split[0],
 			Pass: split[1],
-		}).AsSha256Mechanism()
+		}
+		saslScram = a.AsSha256Mechanism()
+		if len(split) == 3 {
+			n, err := strconv.Atoi(split[2])
+			if err != nil {
+				panic(fmt.Sprintf("invalid scram alg %q: %v", split[2], err))
+			}
+			if n != 256 && n != 512 {
+				panic(fmt.Sprintf("invalid scram alg %q: must be 256 or 512", split[2]))
+			}
+			if n == 512 {
+				saslScram = a.AsSha512Mechanism()
+			}
+		}
 	}
 }