mirror of
https://github.com/Significant-Gravitas/AutoGPT.git
synced 2026-02-08 22:05:08 -05:00
88f711e486
To maintain the transaction running-balance integrity, the monthly top-up balance has to adjust the top-up amount to maintain the balance of the user to be at least X amount, instead of top-up the user balance with a fixed X amount. ### Changes 🏗️ * Add an additional query to sum the remaining un-snapshotted balance (if any). * Monthly top-up transaction set the amount target_amount - current_balance instead of fixed target_amount. ### Checklist 📋 #### For code changes: - [ ] I have clearly listed my changes in the PR description - [ ] I have made a test plan - [ ] I have tested my changes according to the test plan: <!-- Put your test plan here: --> - [ ] ... <details> <summary>Example test plan</summary> - [ ] Create from scratch and execute an agent with at least 3 blocks - [ ] Import an agent from file upload, and confirm it executes correctly - [ ] Upload agent to marketplace - [ ] Import an agent from marketplace and confirm it executes correctly - [ ] Edit an agent from monitor, and confirm it executes correctly </details> #### For configuration changes: - [ ] `.env.example` is updated or already compatible with my changes - [ ] `docker-compose.yml` is updated or already compatible with my changes - [ ] I have included a list of my configuration changes in the PR description (under **Changes**) <details> <summary>Examples of configuration changes</summary> - Changing ports - Adding new services that need to communicate with each other - Secrets or environment variable changes - New or infrastructure changes such as databases </details>
105 lines
3.6 KiB
Python
105 lines
3.6 KiB
Python
from datetime import datetime, timezone
|
|
|
|
import pytest
|
|
from prisma.models import CreditTransaction
|
|
|
|
from backend.blocks.llm import AITextGeneratorBlock
|
|
from backend.data.credit import BetaUserCredit
|
|
from backend.data.execution import NodeExecutionEntry
|
|
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 = BetaUserCredit(REFILL_VALUE)
|
|
|
|
|
|
async def disable_test_user_transactions():
|
|
await CreditTransaction.prisma().delete_many(where={"userId": DEFAULT_USER_ID})
|
|
|
|
|
|
@pytest.mark.asyncio(scope="session")
|
|
async def test_block_credit_usage(server: SpinTestServer):
|
|
await disable_test_user_transactions()
|
|
await user_credit.top_up_credits(DEFAULT_USER_ID, 100)
|
|
current_credit = await user_credit.get_credits(DEFAULT_USER_ID)
|
|
|
|
spending_amount_1 = await user_credit.spend_credits(
|
|
NodeExecutionEntry(
|
|
user_id=DEFAULT_USER_ID,
|
|
graph_id="test_graph",
|
|
node_id="test_node",
|
|
graph_exec_id="test_graph_exec",
|
|
node_exec_id="test_node_exec",
|
|
block_id=AITextGeneratorBlock().id,
|
|
data={
|
|
"model": "gpt-4-turbo",
|
|
"credentials": {
|
|
"id": openai_credentials.id,
|
|
"provider": openai_credentials.provider,
|
|
"type": openai_credentials.type,
|
|
},
|
|
},
|
|
),
|
|
0.0,
|
|
0.0,
|
|
)
|
|
assert spending_amount_1 > 0
|
|
|
|
spending_amount_2 = await user_credit.spend_credits(
|
|
NodeExecutionEntry(
|
|
user_id=DEFAULT_USER_ID,
|
|
graph_id="test_graph",
|
|
node_id="test_node",
|
|
graph_exec_id="test_graph_exec",
|
|
node_exec_id="test_node_exec",
|
|
block_id=AITextGeneratorBlock().id,
|
|
data={"model": "gpt-4-turbo", "api_key": "owned_api_key"},
|
|
),
|
|
0.0,
|
|
0.0,
|
|
)
|
|
assert spending_amount_2 == 0
|
|
|
|
new_credit = await user_credit.get_credits(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):
|
|
await disable_test_user_transactions()
|
|
current_credit = await user_credit.get_credits(DEFAULT_USER_ID)
|
|
|
|
await user_credit.top_up_credits(DEFAULT_USER_ID, 100)
|
|
|
|
new_credit = await user_credit.get_credits(DEFAULT_USER_ID)
|
|
assert new_credit == current_credit + 100
|
|
|
|
|
|
@pytest.mark.asyncio(scope="session")
|
|
async def test_block_credit_reset(server: SpinTestServer):
|
|
await disable_test_user_transactions()
|
|
month1 = 1
|
|
month2 = 2
|
|
|
|
# set the calendar to month 2 but use current time from now
|
|
user_credit.time_now = lambda: datetime.now(timezone.utc).replace(month=month2)
|
|
month2credit = await user_credit.get_credits(DEFAULT_USER_ID)
|
|
|
|
# Month 1 result should only affect month 1
|
|
user_credit.time_now = lambda: datetime.now(timezone.utc).replace(month=month1)
|
|
month1credit = await user_credit.get_credits(DEFAULT_USER_ID)
|
|
await user_credit.top_up_credits(DEFAULT_USER_ID, 100)
|
|
assert await user_credit.get_credits(DEFAULT_USER_ID) == month1credit + 100
|
|
|
|
# Month 2 balance is unaffected
|
|
user_credit.time_now = lambda: datetime.now(timezone.utc).replace(month=month2)
|
|
assert await user_credit.get_credits(DEFAULT_USER_ID) == month2credit
|
|
|
|
|
|
@pytest.mark.asyncio(scope="session")
|
|
async def test_credit_refill(server: SpinTestServer):
|
|
await disable_test_user_transactions()
|
|
balance = await user_credit.get_credits(DEFAULT_USER_ID)
|
|
assert balance == REFILL_VALUE
|