From 581f0741600ba518a5a066cc7fd279da9e463c1c Mon Sep 17 00:00:00 2001
From: agriffaut <agriffaut@gmail.com>
Date: Thu, 30 Dec 2021 10:31:13 +0100
Subject: [PATCH] Add functional test for admin client quotas

---
 functional_admin_test.go | 109 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 109 insertions(+)
 create mode 100644 functional_admin_test.go

diff --git a/functional_admin_test.go b/functional_admin_test.go
new file mode 100644
index 0000000000..c20694bcd6
--- /dev/null
+++ b/functional_admin_test.go
@@ -0,0 +1,109 @@
+//go:build functional
+// +build functional
+
+package sarama
+
+import (
+	"testing"
+)
+
+func TestFuncAdminQuotas(t *testing.T) {
+	checkKafkaVersion(t, "2.6.0.0")
+	setupFunctionalTest(t)
+	defer teardownFunctionalTest(t)
+
+	kafkaVersion, err := ParseKafkaVersion(FunctionalTestEnv.KafkaVersion)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	config := NewTestConfig()
+	config.Version = kafkaVersion
+	adminClient, err := NewClusterAdmin(FunctionalTestEnv.KafkaBrokerAddrs, config)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Check that we can read the quotas, and that they are empty
+	quotas, err := adminClient.DescribeClientQuotas(nil, false)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(quotas) != 0 {
+		t.Errorf("Expected quotas to be empty at start, found: %v", quotas)
+	}
+
+	// Put a quota on default user
+	// /config/users/<default>
+	component := QuotaEntityComponent{
+		EntityType: QuotaEntityUser,
+		MatchType:  QuotaMatchDefault,
+	}
+	op := ClientQuotasOp{
+		Key:   "producer_byte_rate",
+		Value: 1024,
+	}
+	err = adminClient.AlterClientQuotas([]QuotaEntityComponent{component}, op, false)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Check that we now have a quota entry
+	quotas, err = adminClient.DescribeClientQuotas(nil, false)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(quotas) == 0 {
+		t.Error("Expected not empty quotas")
+	}
+	if len(quotas) > 1 {
+		t.Errorf("Expected one quota entry, found: %v", quotas)
+	}
+
+	// Put a quota on specific client-id for a specific user
+	// /config/users/<user>/clients/<client-id>
+	userEntityComponent := QuotaEntityComponent{
+		EntityType: QuotaEntityUser,
+		MatchType:  QuotaMatchExact,
+		Name:       "sarama",
+	}
+	clientEntityComponent := QuotaEntityComponent{
+		EntityType: QuotaEntityClientID,
+		MatchType:  QuotaMatchExact,
+		Name:       "sarama-consumer",
+	}
+	op = ClientQuotasOp{
+		Key:   "consumer_byte_rate",
+		Value: 2048,
+	}
+
+	err = adminClient.AlterClientQuotas([]QuotaEntityComponent{userEntityComponent, clientEntityComponent}, op, false)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Check that we can query a specific quota entry
+	userFilterComponent := QuotaFilterComponent{
+		EntityType: QuotaEntityUser,
+		MatchType:  QuotaMatchExact,
+		Match:      "sarama",
+	}
+	clientFilterComponent := QuotaFilterComponent{
+		EntityType: QuotaEntityClientID,
+		MatchType:  QuotaMatchExact,
+		Match:      "sarama-consumer",
+	}
+	quotas, err = adminClient.DescribeClientQuotas([]QuotaFilterComponent{userFilterComponent, clientFilterComponent}, true)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(quotas) == 0 {
+		t.Error("Expected not empty quotas")
+	}
+	if len(quotas) > 1 {
+		t.Errorf("Expected one quota entry, found: %v", quotas)
+	}
+	if quotas[0].Values[op.Key] != op.Value {
+		t.Errorf("Expected specific quota value to be %f, found: %v", op.Value, quotas[0].Values[op.Key])
+	}
+}