From e8ba3ee37e3edc925b16a7fb7f11f68b2f9c31e4 Mon Sep 17 00:00:00 2001 From: Shivam Mishra Date: Fri, 17 Jan 2025 16:10:14 +0530 Subject: [PATCH] feat: add sepcs for captain usage --- enterprise/app/models/enterprise/account.rb | 40 ++++++------ spec/enterprise/models/account_spec.rb | 72 ++++++++++++++++++++- 2 files changed, 92 insertions(+), 20 deletions(-) diff --git a/enterprise/app/models/enterprise/account.rb b/enterprise/app/models/enterprise/account.rb index f8450b12d..03678d7a8 100644 --- a/enterprise/app/models/enterprise/account.rb +++ b/enterprise/app/models/enterprise/account.rb @@ -1,5 +1,5 @@ module Enterprise::Account - ACCOUNT_RESPONSES_LIMIT = 'ACCOUNT_RESPONSES_LIMIT'.freeze + CAPTAIN_RESPONSES = 'captain_responses'.freeze def usage_limits { @@ -12,6 +12,17 @@ module Enterprise::Account } end + def increment_response_usage + current_usage = self[:limits][CAPTAIN_RESPONSES].to_i || 0 + self[:limits][CAPTAIN_RESPONSES] = current_usage + 1 + save + end + + def reset_response_usage + self[:limits][CAPTAIN_RESPONSES] = 0 + save + end + def subscribed_features plan_features = InstallationConfig.find_by(name: 'CHATWOOT_CLOUD_PLAN_FEATURES')&.value return [] if plan_features.blank? @@ -30,34 +41,24 @@ module Enterprise::Account private def get_captain_limits(type) - total_count = captain_monthly_limit[type].to_i + total_count = captain_monthly_limit[type.to_s].to_i consumed = if type == :documents - captain_documents.available.count + captain_documents.count else - self[:limits][ACCOUNT_RESPONSES_LIMIT] || 0 + self[:limits][CAPTAIN_RESPONSES] || 0 end + consumed = 0 if consumed.negative? + { total_count: total_count, - current_available: [total_count - consumed, 0].max, + # ensure current_available is never negative or more than total_count + current_available: [[total_count - consumed, 0].max, total_count].min, consumed: consumed } end - def increment_response_usage - config_name = ACCOUNT_RESPONSES_LIMIT - self[:limits][config_name] = 0 if self[:limits][config_name].blank? - self[:limits][config_name] = self[:limits][config_name].to_i + 1 - save - end - - def reset_response_usage - config_name = ACCOUNT_RESPONSES_LIMIT - self[:limits][config_name] = 0 - save - end - def plan_name custom_attributes['plan_name'] end @@ -84,7 +85,8 @@ module Enterprise::Account 'type' => 'object', 'properties' => { 'inboxes' => { 'type': 'number' }, - 'agents' => { 'type': 'number' } + 'agents' => { 'type': 'number' }, + 'captain_responses' => { 'type': 'number' } }, 'required' => [], 'additionalProperties' => false diff --git a/spec/enterprise/models/account_spec.rb b/spec/enterprise/models/account_spec.rb index 8e2681a9b..51f7e3bc6 100644 --- a/spec/enterprise/models/account_spec.rb +++ b/spec/enterprise/models/account_spec.rb @@ -28,11 +28,81 @@ RSpec.describe Account, type: :model do end describe 'usage_limits' do + let(:captain_limits) do + { + :startups => { :documents => 100, :responses => 100 }, + :business => { :documents => 200, :responses => 300 }, + :enterprise => { :documents => 300, :responses => 500 } + }.with_indifferent_access + end + let!(:account) { create(:account, { custom_attributes: { plan_name: 'startups' } }) } + let!(:assistant) { create(:captain_assistant, account: account) } + let!(:document) { create(:captain_document, assistant: assistant, account: account, status: :available) } + before do create(:installation_config, name: 'ACCOUNT_AGENTS_LIMIT', value: 20) + create(:installation_config, name: 'CAPTAIN_CLOUD_PLAN_LIMITS', value: captain_limits.to_json) end - let!(:account) { create(:account) } + describe 'captain limits' do + it 'returns monthly limit accurately' do + %w[startups business enterprise].each do |plan| + account.custom_attributes = { 'plan_name': plan } + account.save! + expect(account.captain_monthly_limit).to eq captain_limits[plan] + end + end + + it 'incrementing responses updates usage_limits' do + account.increment_response_usage + + responses_limits = account.usage_limits[:captain][:generated_responses] + + expect(account.limits['captain_responses']).to eq 1 + expect(responses_limits[:consumed]).to eq 1 + expect(responses_limits[:current_available]).to eq captain_limits[:startups][:responses] - 1 + end + + it 'reseting responses limits updates usage_limits' do + account.limits['captain_responses'] = 30 + account.save! + + responses_limits = account.usage_limits[:captain][:generated_responses] + + expect(responses_limits[:consumed]).to eq 30 + expect(responses_limits[:current_available]).to eq captain_limits[:startups][:responses] - 30 + + account.reset_response_usage + responses_limits = account.usage_limits[:captain][:generated_responses] + + expect(account.limits['captain_responses']).to eq 0 + expect(responses_limits[:consumed]).to eq 0 + expect(responses_limits[:current_available]).to eq captain_limits[:startups][:responses] + end + + it 'reflects document limits' do + document_limits = account.usage_limits[:captain][:documents] + + expect(document_limits[:consumed]).to eq 1 + expect(document_limits[:current_available]).to eq captain_limits[:startups][:documents] - 1 + end + + it 'current_available is never out of bounds' do + account.limits['captain_responses'] = 3000 + account.save! + + responses_limits = account.usage_limits[:captain][:generated_responses] + expect(responses_limits[:consumed]).to eq 3000 + expect(responses_limits[:current_available]).to eq 0 + + account.limits['captain_responses'] = -100 + account.save! + + responses_limits = account.usage_limits[:captain][:generated_responses] + expect(responses_limits[:consumed]).to eq 0 + expect(responses_limits[:current_available]).to eq captain_limits[:startups][:responses] + end + end describe 'audit logs' do it 'returns audit logs' do