From 8e50d59a445536df7f56249d8411d6f05d953990 Mon Sep 17 00:00:00 2001
From: iliax <ikuramshin@provectus.com>
Date: Fri, 11 Aug 2023 18:56:45 +0400
Subject: [PATCH] wip

---
 .../ui/service/quota/ClientQuotaService.java  | 11 +++--
 .../service/quota/ClientQuotaServiceTest.java | 44 +++++++++++++++++++
 .../main/resources/swagger/kafka-ui-api.yaml  |  1 +
 3 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/kafka-ui-api/src/main/java/com/provectus/kafka/ui/service/quota/ClientQuotaService.java b/kafka-ui-api/src/main/java/com/provectus/kafka/ui/service/quota/ClientQuotaService.java
index f132cf17ea..3c78a9bf81 100644
--- a/kafka-ui-api/src/main/java/com/provectus/kafka/ui/service/quota/ClientQuotaService.java
+++ b/kafka-ui-api/src/main/java/com/provectus/kafka/ui/service/quota/ClientQuotaService.java
@@ -4,6 +4,7 @@ import static org.apache.kafka.common.quota.ClientQuotaEntity.CLIENT_ID;
 import static org.apache.kafka.common.quota.ClientQuotaEntity.IP;
 import static org.apache.kafka.common.quota.ClientQuotaEntity.USER;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.Sets;
 import com.provectus.kafka.ui.exception.ValidationException;
 import com.provectus.kafka.ui.model.KafkaCluster;
@@ -38,8 +39,11 @@ public class ClientQuotaService {
     return adminClientService.get(cluster)
         .flatMap(ac -> ac.getClientQuotas(ClientQuotaFilter.all()))
         .flatMapIterable(map ->
-            map.entrySet().stream().map(e -> ClientQuotaRecord.create(e.getKey(), e.getValue())).toList())
-        .sort(ClientQuotaRecord.COMPARATOR);
+            map.entrySet().stream()
+                .map(e -> ClientQuotaRecord.create(e.getKey(), e.getValue()))
+                .sorted(ClientQuotaRecord.COMPARATOR)
+                .toList()
+        );
   }
 
   //returns 201 if new entity was created, 200 if existing was updated, 204 if existing was deleted
@@ -66,7 +70,8 @@ public class ClientQuotaService {
         );
   }
 
-  private ClientQuotaEntity quotaEntity(@Nullable String user, @Nullable String clientId, @Nullable String ip) {
+  @VisibleForTesting
+  static ClientQuotaEntity quotaEntity(@Nullable String user, @Nullable String clientId, @Nullable String ip) {
     if (Stream.of(user, clientId, ip).allMatch(Objects::isNull)) {
       throw new ValidationException("Quota entity id is not set");
     }
diff --git a/kafka-ui-api/src/test/java/com/provectus/kafka/ui/service/quota/ClientQuotaServiceTest.java b/kafka-ui-api/src/test/java/com/provectus/kafka/ui/service/quota/ClientQuotaServiceTest.java
index 1ca42f9de8..d79a62c78a 100644
--- a/kafka-ui-api/src/test/java/com/provectus/kafka/ui/service/quota/ClientQuotaServiceTest.java
+++ b/kafka-ui-api/src/test/java/com/provectus/kafka/ui/service/quota/ClientQuotaServiceTest.java
@@ -1,11 +1,55 @@
 package com.provectus.kafka.ui.service.quota;
 
+import static org.assertj.core.api.Assertions.assertThat;
+
 import com.provectus.kafka.ui.AbstractIntegrationTest;
+import com.provectus.kafka.ui.model.KafkaCluster;
+import com.provectus.kafka.ui.service.ClustersStorage;
+import java.util.Map;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
 import org.springframework.beans.factory.annotation.Autowired;
+import reactor.test.StepVerifier;
 
 class ClientQuotaServiceTest extends AbstractIntegrationTest {
 
   @Autowired
   ClientQuotaService quotaService;
 
+  private KafkaCluster cluster;
+
+  @BeforeEach
+  void init() {
+    cluster = applicationContext.getBean(ClustersStorage.class).getClusterByName(LOCAL).get();
+  }
+
+  @ParameterizedTest
+  @CsvSource(
+      value = {
+          "testUser, null, null ",
+          "null, testUserId, null",
+          "testUser2, testUserId2, null",
+          "null, null, 127.0.0.1"
+      },
+      nullValues = "null"
+  )
+  void createsQuotaRecord(String user, String clientId, String ip) {
+    StepVerifier.create(
+            quotaService.upsert(
+                cluster,
+                user,
+                clientId,
+                ip,
+                Map.of(
+                    "producer_byte_rate", 123.0,
+                    "consumer_byte_rate", 234.0,
+                    "request_percentage", 10.0
+                )
+            )
+        )
+        .assertNext(status -> assertThat(status.value()).isEqualTo(201))
+        .verifyComplete();
+  }
+
 }
diff --git a/kafka-ui-contract/src/main/resources/swagger/kafka-ui-api.yaml b/kafka-ui-contract/src/main/resources/swagger/kafka-ui-api.yaml
index 57a0550616..98fe8eed2e 100644
--- a/kafka-ui-contract/src/main/resources/swagger/kafka-ui-api.yaml
+++ b/kafka-ui-contract/src/main/resources/swagger/kafka-ui-api.yaml
@@ -3650,6 +3650,7 @@ components:
         - KSQL
         - ACL
         - AUDIT
+        - CLIENT_QUOTAS
 
     KafkaAcl:
       type: object