Back to Repositories

Validating User Credit Management System in AutoGPT

A comprehensive test suite for validating user credit management functionality in AutoGPT, focusing on credit usage tracking, top-up operations, and monthly credit reset mechanisms. These tests ensure proper handling of user credits across different scenarios and time periods.

Test Coverage Overview

The test suite provides extensive coverage of credit management operations:

Credit usage validation for AI text generation blocks
Credit balance updates and top-up functionality
Monthly credit reset mechanisms
Credit refill operations

Edge cases include different API key scenarios and cross-month credit management.

Implementation Analysis

Tests utilize pytest’s async testing capabilities with session-scoped fixtures. The implementation follows a structured approach using mock time functions and database operations through Prisma ORM. Framework-specific features include pytest.mark.asyncio decorators and SpinTestServer integration.

Technical Details

Testing Framework: pytest with async support
Database: Prisma ORM for credit transactions
Custom Utilities: SpinTestServer for test environment
Mocking: Custom time_now lambda for date manipulation
Test Scope: Session-level async tests

Best Practices Demonstrated

The test suite exemplifies robust testing practices with isolated test cases, proper setup/teardown management, and comprehensive assertion coverage. Notable practices include transaction isolation, time-based testing patterns, and clear separation of test scenarios with focused assertions for each use case.

significant-gravitas/autogpt

autogpt_platform/backend/test/data/test_credit.py

            
from datetime import datetime

import pytest
from prisma.models import UserBlockCredit

from backend.blocks.llm import AITextGeneratorBlock
from backend.data.credit import UserCredit
from backend.data.user import DEFAULT_USER_ID
from backend.integrations.credentials_store import openai_credentials
from backend.util.test import SpinTestServer

REFILL_VALUE = 1000
user_credit = UserCredit(REFILL_VALUE)


@pytest.mark.asyncio(scope="session")
async def test_block_credit_usage(server: SpinTestServer):
    current_credit = await user_credit.get_or_refill_credit(DEFAULT_USER_ID)

    spending_amount_1 = await user_credit.spend_credits(
        DEFAULT_USER_ID,
        current_credit,
        AITextGeneratorBlock().id,
        {
            "model": "gpt-4-turbo",
            "credentials": {
                "id": openai_credentials.id,
                "provider": openai_credentials.provider,
                "type": openai_credentials.type,
            },
        },
        0.0,
        0.0,
        validate_balance=False,
    )
    assert spending_amount_1 > 0

    spending_amount_2 = await user_credit.spend_credits(
        DEFAULT_USER_ID,
        current_credit,
        AITextGeneratorBlock().id,
        {"model": "gpt-4-turbo", "api_key": "owned_api_key"},
        0.0,
        0.0,
        validate_balance=False,
    )
    assert spending_amount_2 == 0

    new_credit = await user_credit.get_or_refill_credit(DEFAULT_USER_ID)
    assert new_credit == current_credit - spending_amount_1 - spending_amount_2


@pytest.mark.asyncio(scope="session")
async def test_block_credit_top_up(server: SpinTestServer):
    current_credit = await user_credit.get_or_refill_credit(DEFAULT_USER_ID)

    await user_credit.top_up_credits(DEFAULT_USER_ID, 100)

    new_credit = await user_credit.get_or_refill_credit(DEFAULT_USER_ID)
    assert new_credit == current_credit + 100


@pytest.mark.asyncio(scope="session")
async def test_block_credit_reset(server: SpinTestServer):
    month1 = datetime(2022, 1, 15)
    month2 = datetime(2022, 2, 15)

    user_credit.time_now = lambda: month2
    month2credit = await user_credit.get_or_refill_credit(DEFAULT_USER_ID)

    # Month 1 result should only affect month 1
    user_credit.time_now = lambda: month1
    month1credit = await user_credit.get_or_refill_credit(DEFAULT_USER_ID)
    await user_credit.top_up_credits(DEFAULT_USER_ID, 100)
    assert await user_credit.get_or_refill_credit(DEFAULT_USER_ID) == month1credit + 100

    # Month 2 balance is unaffected
    user_credit.time_now = lambda: month2
    assert await user_credit.get_or_refill_credit(DEFAULT_USER_ID) == month2credit


@pytest.mark.asyncio(scope="session")
async def test_credit_refill(server: SpinTestServer):
    # Clear all transactions within the month
    await UserBlockCredit.prisma().update_many(
        where={
            "userId": DEFAULT_USER_ID,
            "createdAt": {
                "gte": datetime(2022, 2, 1),
                "lt": datetime(2022, 3, 1),
            },
        },
        data={"isActive": False},
    )
    user_credit.time_now = lambda: datetime(2022, 2, 15)

    balance = await user_credit.get_or_refill_credit(DEFAULT_USER_ID)
    assert balance == REFILL_VALUE