From 87ee07ea6efb3dc9708c453faadd39ec9c7586ab Mon Sep 17 00:00:00 2001
From: Savindu Dimal <dimalsavindu@gmail.com>
Date: Mon, 14 Oct 2024 15:37:39 +0530
Subject: [PATCH] Fix issue: token based throttling can not be removed

---
 .../utils/RegistryPersistenceUtil.java        | 36 +++++++------------
 1 file changed, 12 insertions(+), 24 deletions(-)

diff --git a/components/apimgt/org.wso2.carbon.apimgt.persistence/src/main/java/org/wso2/carbon/apimgt/persistence/utils/RegistryPersistenceUtil.java b/components/apimgt/org.wso2.carbon.apimgt.persistence/src/main/java/org/wso2/carbon/apimgt/persistence/utils/RegistryPersistenceUtil.java
index a7a4658da4da..3cc2e386591a 100644
--- a/components/apimgt/org.wso2.carbon.apimgt.persistence/src/main/java/org/wso2/carbon/apimgt/persistence/utils/RegistryPersistenceUtil.java
+++ b/components/apimgt/org.wso2.carbon.apimgt.persistence/src/main/java/org/wso2/carbon/apimgt/persistence/utils/RegistryPersistenceUtil.java
@@ -222,30 +222,18 @@ public static GenericArtifact createAPIArtifactContent(GenericArtifact artifact,
                         .getTokenBasedThrottlingConfiguration();
                 artifact.setAttribute(APIConstants.API_TOKEN_BASED_THROTTLING_ENABLED,
                         tokenBasedThrottlingCountHolder.isTokenBasedThrottlingEnabled().toString());
-                if (tokenBasedThrottlingCountHolder.getProductionMaxPromptTokenCount() != null) {
-                    artifact.setAttribute(APIConstants.API_PRODUCTION_MAX_PROMPT_TOKEN_COUNT,
-                            tokenBasedThrottlingCountHolder.getProductionMaxPromptTokenCount());
-                }
-                if (tokenBasedThrottlingCountHolder.getProductionMaxCompletionTokenCount() != null) {
-                    artifact.setAttribute(APIConstants.API_PRODUCTION_MAX_COMPLETION_TOKEN_COUNT,
-                            tokenBasedThrottlingCountHolder.getProductionMaxCompletionTokenCount());
-                }
-                if (tokenBasedThrottlingCountHolder.getProductionMaxTotalTokenCount() != null) {
-                    artifact.setAttribute(APIConstants.API_PRODUCTION_MAX_TOTAL_TOKEN_COUNT,
-                            tokenBasedThrottlingCountHolder.getProductionMaxTotalTokenCount());
-                }
-                if (tokenBasedThrottlingCountHolder.getSandboxMaxPromptTokenCount() != null) {
-                    artifact.setAttribute(APIConstants.API_SANDBOX_MAX_PROMPT_TOKEN_COUNT,
-                            tokenBasedThrottlingCountHolder.getSandboxMaxPromptTokenCount());
-                }
-                if (tokenBasedThrottlingCountHolder.getSandboxMaxCompletionTokenCount() != null) {
-                    artifact.setAttribute(APIConstants.API_SANDBOX_MAX_COMPLETION_TOKEN_COUNT,
-                            tokenBasedThrottlingCountHolder.getSandboxMaxCompletionTokenCount());
-                }
-                if (tokenBasedThrottlingCountHolder.getSandboxMaxTotalTokenCount() != null) {
-                    artifact.setAttribute(APIConstants.API_SANDBOX_MAX_TOTAL_TOKEN_COUNT,
-                            tokenBasedThrottlingCountHolder.getSandboxMaxTotalTokenCount());
-                }
+                artifact.setAttribute(APIConstants.API_PRODUCTION_MAX_PROMPT_TOKEN_COUNT,
+                        tokenBasedThrottlingCountHolder.getProductionMaxPromptTokenCount());
+                artifact.setAttribute(APIConstants.API_PRODUCTION_MAX_COMPLETION_TOKEN_COUNT,
+                        tokenBasedThrottlingCountHolder.getProductionMaxCompletionTokenCount());
+                artifact.setAttribute(APIConstants.API_PRODUCTION_MAX_TOTAL_TOKEN_COUNT,
+                        tokenBasedThrottlingCountHolder.getProductionMaxTotalTokenCount());
+                artifact.setAttribute(APIConstants.API_SANDBOX_MAX_PROMPT_TOKEN_COUNT,
+                        tokenBasedThrottlingCountHolder.getSandboxMaxPromptTokenCount());
+                artifact.setAttribute(APIConstants.API_SANDBOX_MAX_COMPLETION_TOKEN_COUNT,
+                        tokenBasedThrottlingCountHolder.getSandboxMaxCompletionTokenCount());
+                artifact.setAttribute(APIConstants.API_SANDBOX_MAX_TOTAL_TOKEN_COUNT,
+                        tokenBasedThrottlingCountHolder.getSandboxMaxTotalTokenCount());
             }
 
             String policies = policyBuilder.toString();