discourse-ai/spec/models/llm_quota_usage_spec.rb

# frozen_string_literal: true

RSpec.describe LlmQuotaUsage do
  fab!(:group)
  fab!(:user)
  fab!(:llm_model)

  fab!(:llm_quota) do
    Fabricate(
      :llm_quota,
      group: group,
      llm_model: llm_model,
      max_tokens: 1000,
      max_usages: 10,
      duration_seconds: 1.day.to_i,
    )
  end

  describe ".find_or_create_for" do
    it "creates a new usage record if none exists" do
      freeze_time

      usage = described_class.find_or_create_for(user: user, llm_quota: llm_quota)

      expect(usage).to be_persisted
      expect(usage.started_at).to eq_time(Time.current)
      expect(usage.reset_at).to eq_time(Time.current + llm_quota.duration_seconds.seconds)
      expect(usage.input_tokens_used).to eq(0)
      expect(usage.output_tokens_used).to eq(0)
      expect(usage.usages).to eq(0)
    end

    it "returns existing usage record if one exists" do
      existing = Fabricate(:llm_quota_usage, user: user, llm_quota: llm_quota)

      usage = described_class.find_or_create_for(user: user, llm_quota: llm_quota)

      expect(usage.id).to eq(existing.id)
    end
  end

  describe "#reset_if_needed!" do
    let(:usage) { Fabricate(:llm_quota_usage, user: user, llm_quota: llm_quota) }

    it "resets usage when past reset_at" do
      usage.update!(
        input_tokens_used: 100,
        output_tokens_used: 200,
        usages: 5,
        reset_at: 1.minute.ago,
      )

      freeze_time

      usage.reset_if_needed!

      expect(usage.reload.input_tokens_used).to eq(0)
      expect(usage.output_tokens_used).to eq(0)
      expect(usage.usages).to eq(0)
      expect(usage.started_at).to eq_time(Time.current)
      expect(usage.reset_at).to eq_time(Time.current + llm_quota.duration_seconds.seconds)
    end

    it "doesn't reset if reset_at hasn't passed" do
      freeze_time

      original_values = {
        input_tokens_used: 100,
        output_tokens_used: 200,
        usages: 5,
        reset_at: 1.minute.from_now,
      }

      usage.update!(original_values)
      usage.reset_if_needed!

      usage.reload
      expect(usage.input_tokens_used).to eq(original_values[:input_tokens_used])
      expect(usage.output_tokens_used).to eq(original_values[:output_tokens_used])
      expect(usage.usages).to eq(original_values[:usages])
      expect(usage.reset_at).to eq_time(original_values[:reset_at])
    end
  end

  describe "#increment_usage!" do
    let(:usage) { Fabricate(:llm_quota_usage, user: user, llm_quota: llm_quota) }

    it "increments usage counts" do
      usage.increment_usage!(input_tokens: 50, output_tokens: 30)

      expect(usage.reload.input_tokens_used).to eq(50)
      expect(usage.output_tokens_used).to eq(30)
      expect(usage.usages).to eq(1)
    end

    it "accumulates multiple increments" do
      2.times { usage.increment_usage!(input_tokens: 50, output_tokens: 30) }

      expect(usage.reload.input_tokens_used).to eq(100)
      expect(usage.output_tokens_used).to eq(60)
      expect(usage.usages).to eq(2)
    end

    it "resets counts if needed before incrementing" do
      usage.update!(
        input_tokens_used: 100,
        output_tokens_used: 200,
        usages: 5,
        reset_at: 1.minute.ago,
      )

      usage.increment_usage!(input_tokens: 50, output_tokens: 30)

      expect(usage.reload.input_tokens_used).to eq(50)
      expect(usage.output_tokens_used).to eq(30)
      expect(usage.usages).to eq(1)
    end
  end

  describe "#check_quota!" do
    let(:usage) { Fabricate(:llm_quota_usage, user: user, llm_quota: llm_quota) }

    it "doesn't raise error when within limits" do
      expect { usage.check_quota! }.not_to raise_error
    end

    it "raises error when max_tokens exceeded" do
      usage.update!(input_tokens_used: llm_quota.max_tokens + 1)

      expect { usage.check_quota! }.to raise_error(LlmQuotaUsage::QuotaExceededError, /exceeded/)
    end

    it "raises error when max_usages exceeded" do
      usage.update!(usages: llm_quota.max_usages + 1)

      expect { usage.check_quota! }.to raise_error(LlmQuotaUsage::QuotaExceededError, /exceeded/)
    end

    it "resets quota if needed before checking" do
      usage.update!(input_tokens_used: llm_quota.max_tokens + 1, reset_at: 1.minute.ago)

      expect { usage.check_quota! }.not_to raise_error
      expect(usage.reload.input_tokens_used).to eq(0)
    end
  end

  describe "#quota_exceeded?" do
    let(:usage) { Fabricate(:llm_quota_usage, user: user, llm_quota: llm_quota) }

    it "returns false when within limits" do
      expect(usage.quota_exceeded?).to be false
    end

    it "returns true when max_tokens exceeded" do
      usage.update!(input_tokens_used: llm_quota.max_tokens + 1)
      expect(usage.quota_exceeded?).to be true
    end

    it "returns true when max_usages exceeded" do
      usage.update!(usages: llm_quota.max_usages + 1)
      expect(usage.quota_exceeded?).to be true
    end

    it "returns false when quota is nil" do
      tokens = llm_quota.max_tokens + 1
      usage.llm_quota.update!(max_tokens: nil)
      usage.update!(input_tokens_used: tokens)
      expect(usage.quota_exceeded?).to be false
    end
  end

  describe "calculation methods" do
    let(:usage) { Fabricate(:llm_quota_usage, user: user, llm_quota: llm_quota) }

    describe "#total_tokens_used" do
      it "sums input and output tokens" do
        usage.update!(input_tokens_used: 100, output_tokens_used: 200)
        expect(usage.total_tokens_used).to eq(300)
      end
    end

    describe "#remaining_tokens" do
      it "calculates remaining tokens when under limit" do
        usage.update!(input_tokens_used: 300, output_tokens_used: 200)
        expect(usage.remaining_tokens).to eq(500)
      end

      it "returns 0 when over limit" do
        usage.update!(input_tokens_used: 800, output_tokens_used: 300)
        expect(usage.remaining_tokens).to eq(0)
      end

      it "returns nil when no max_tokens set" do
        usage.llm_quota.update!(max_tokens: nil)
        expect(usage.remaining_tokens).to be_nil
      end
    end

    describe "#remaining_usages" do
      it "calculates remaining usages when under limit" do
        usage.update!(usages: 7)
        expect(usage.remaining_usages).to eq(3)
      end

      it "returns 0 when over limit" do
        usage.update!(usages: 15)
        expect(usage.remaining_usages).to eq(0)
      end

      it "returns nil when no max_usages set" do
        usage.llm_quota.update!(max_usages: nil)
        expect(usage.remaining_usages).to be_nil
      end
    end

    describe "#percentage_tokens_used" do
      it "calculates percentage correctly" do
        usage.update!(input_tokens_used: 250, output_tokens_used: 250)
        expect(usage.percentage_tokens_used).to eq(50)
      end

      it "caps at 100%" do
        usage.update!(input_tokens_used: 2000)
        expect(usage.percentage_tokens_used).to eq(100)
      end

      it "returns 0 when no max_tokens set" do
        usage.llm_quota.update!(max_tokens: nil)
        expect(usage.percentage_tokens_used).to eq(0)
      end
    end

    describe "#percentage_usages_used" do
      it "calculates percentage correctly" do
        usage.update!(usages: 5)
        expect(usage.percentage_usages_used).to eq(50)
      end

      it "caps at 100%" do
        usage.update!(usages: 20)
        expect(usage.percentage_usages_used).to eq(100)
      end

      it "returns 0 when no max_usages set" do
        usage.llm_quota.update!(max_usages: nil)
        expect(usage.percentage_usages_used).to eq(0)
      end
    end
  end
end