From d2ca5f43908584a35e1a221f60abf8208fab6fb6 Mon Sep 17 00:00:00 2001
From: Shivam Mishra <scm.mymail@gmail.com>
Date: Tue, 21 Jan 2025 14:24:50 +0530
Subject: [PATCH] feat: allow limits from custom attributes

---
 enterprise/app/models/enterprise/account.rb | 13 +++++++++++--
 spec/enterprise/models/account_spec.rb      |  2 +-
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/enterprise/app/models/enterprise/account.rb b/enterprise/app/models/enterprise/account.rb
index cde366f1a..239bd0ac8 100644
--- a/enterprise/app/models/enterprise/account.rb
+++ b/enterprise/app/models/enterprise/account.rb
@@ -37,8 +37,8 @@ module Enterprise::Account
     plan_features[plan_name]
   end
 
-  def captain_monthly_limit
-    default_value = { documents: ChatwootApp.max_limit, generated_responses: ChatwootApp.max_limit }.with_indifferent_access
+  def default_captain_limit
+    default_value = { documents: ChatwootApp.max_limit, responses: ChatwootApp.max_limit }.with_indifferent_access
     return default_value if plan_name.blank?
 
     plan_quota = InstallationConfig.find_by(name: 'CAPTAIN_CLOUD_PLAN_LIMITS')&.value
@@ -48,6 +48,15 @@ module Enterprise::Account
     plan_quota[plan_name.downcase]
   end
 
+  def captain_monthly_limit
+    default_limits = default_captain_limit
+
+    {
+      documents: custom_attributes['captain_document_limit'] || default_limits['documents'],
+      responses: custom_attributes['captain_response_limit'] || default_limits['responses']
+    }.with_indifferent_access
+  end
+
   private
 
   def get_captain_limits(type)
diff --git a/spec/enterprise/models/account_spec.rb b/spec/enterprise/models/account_spec.rb
index def07985e..3667c7095 100644
--- a/spec/enterprise/models/account_spec.rb
+++ b/spec/enterprise/models/account_spec.rb
@@ -124,7 +124,7 @@ RSpec.describe Account, type: :model do
       it 'returns default values' do
         account.custom_attributes = { 'plan_name': 'unknown' }
         expect(account.captain_monthly_limit).to eq(
-          { documents: ChatwootApp.max_limit, generated_responses: ChatwootApp.max_limit }.with_indifferent_access
+          { documents: ChatwootApp.max_limit, responses: ChatwootApp.max_limit }.with_indifferent_access
         )
       end
     end