diff --git a/autogpt_platform/backend/backend/api/features/chat/tools/_test_data.py b/autogpt_platform/backend/backend/api/features/chat/tools/_test_data.py index 547dc403d7..507c366b7e 100644 --- a/autogpt_platform/backend/backend/api/features/chat/tools/_test_data.py +++ b/autogpt_platform/backend/backend/api/features/chat/tools/_test_data.py @@ -3,9 +3,8 @@ from datetime import UTC, datetime from os import getenv import pytest -from pydantic import SecretStr - from prisma.types import ProfileCreateInput +from pydantic import SecretStr from backend.api.features.chat.model import ChatSession from backend.api.features.store import db as store_db diff --git a/autogpt_platform/backend/backend/data/credit_ceiling_test.py b/autogpt_platform/backend/backend/data/credit_ceiling_test.py index 6179581338..e4d90d51de 100644 --- a/autogpt_platform/backend/backend/data/credit_ceiling_test.py +++ b/autogpt_platform/backend/backend/data/credit_ceiling_test.py @@ -11,11 +11,7 @@ import pytest from prisma.enums import CreditTransactionType from prisma.errors import UniqueViolationError from prisma.models import CreditTransaction, User, UserBalance -from prisma.types import ( - UserBalanceCreateInput, - UserBalanceUpsertInput, - UserCreateInput, -) +from prisma.types import UserBalanceCreateInput, UserBalanceUpsertInput, UserCreateInput from backend.data.credit import UserCredit from backend.util.json import SafeJson @@ -115,15 +111,15 @@ async def test_ceiling_balance_clamps_when_would_exceed(server: SpinTestServer): ) # Balance should be clamped to ceiling - assert final_balance == 1000, ( - f"Balance should be clamped to 1000, got {final_balance}" - ) + assert ( + final_balance == 1000 + ), f"Balance should be clamped to 1000, got {final_balance}" # Verify with get_credits too stored_balance = await credit_system.get_credits(user_id) - assert stored_balance == 1000, ( - f"Stored balance should be 1000, got {stored_balance}" - ) + assert ( + stored_balance == 1000 + ), f"Stored balance should be 1000, got {stored_balance}" # Verify transaction shows the clamped amount transactions = await CreditTransaction.prisma().find_many( @@ -172,9 +168,9 @@ async def test_ceiling_balance_allows_when_under_threshold(server: SpinTestServe # Verify with get_credits too stored_balance = await credit_system.get_credits(user_id) - assert stored_balance == 500, ( - f"Stored balance should be 500, got {stored_balance}" - ) + assert ( + stored_balance == 500 + ), f"Stored balance should be 500, got {stored_balance}" finally: await cleanup_test_user(user_id) diff --git a/autogpt_platform/backend/backend/data/credit_concurrency_test.py b/autogpt_platform/backend/backend/data/credit_concurrency_test.py index 8df624b47e..965057ed8d 100644 --- a/autogpt_platform/backend/backend/data/credit_concurrency_test.py +++ b/autogpt_platform/backend/backend/data/credit_concurrency_test.py @@ -14,11 +14,7 @@ import pytest from prisma.enums import CreditTransactionType from prisma.errors import UniqueViolationError from prisma.models import CreditTransaction, User, UserBalance -from prisma.types import ( - UserBalanceCreateInput, - UserBalanceUpsertInput, - UserCreateInput, -) +from prisma.types import UserBalanceCreateInput, UserBalanceUpsertInput, UserCreateInput from backend.data.credit import POSTGRES_INT_MAX, UsageTransactionMetadata, UserCredit from backend.util.exceptions import InsufficientBalanceError @@ -116,9 +112,9 @@ async def test_concurrent_spends_same_user(server: SpinTestServer): transactions = await CreditTransaction.prisma().find_many( where={"userId": user_id, "type": prisma.enums.CreditTransactionType.USAGE} ) - assert len(transactions) == 10, ( - f"Expected 10 transactions, got {len(transactions)}" - ) + assert ( + len(transactions) == 10 + ), f"Expected 10 transactions, got {len(transactions)}" finally: await cleanup_test_user(user_id) @@ -329,9 +325,9 @@ async def test_onboarding_reward_idempotency(server: SpinTestServer): "transactionKey": f"REWARD-{user_id}-WELCOME", } ) - assert len(transactions) == 1, ( - f"Expected 1 reward transaction, got {len(transactions)}" - ) + assert ( + len(transactions) == 1 + ), f"Expected 1 reward transaction, got {len(transactions)}" finally: await cleanup_test_user(user_id) @@ -366,9 +362,9 @@ async def test_integer_overflow_protection(server: SpinTestServer): # Balance should be clamped to max_int, not overflowed final_balance = await credit_system.get_credits(user_id) - assert final_balance == max_int, ( - f"Balance should be clamped to {max_int}, got {final_balance}" - ) + assert ( + final_balance == max_int + ), f"Balance should be clamped to {max_int}, got {final_balance}" # Verify transaction was created with clamped amount transactions = await CreditTransaction.prisma().find_many( @@ -379,9 +375,9 @@ async def test_integer_overflow_protection(server: SpinTestServer): order={"createdAt": "desc"}, ) assert len(transactions) > 0, "Transaction should be created" - assert transactions[0].runningBalance == max_int, ( - "Transaction should show clamped balance" - ) + assert ( + transactions[0].runningBalance == max_int + ), "Transaction should show clamped balance" finally: await cleanup_test_user(user_id) @@ -440,9 +436,9 @@ async def test_high_concurrency_stress(server: SpinTestServer): # Verify final balance final_balance = await credit_system.get_credits(user_id) - assert final_balance == expected_balance, ( - f"Expected {expected_balance}, got {final_balance}" - ) + assert ( + final_balance == expected_balance + ), f"Expected {expected_balance}, got {final_balance}" assert final_balance >= 0, "Balance went negative!" finally: @@ -541,9 +537,9 @@ async def test_concurrent_multiple_spends_sufficient_balance(server: SpinTestSer print(f"Successful: {len(successful)}, Failed: {len(failed)}") # All should succeed since 150 - (10 + 20 + 30) = 90 > 0 - assert len(successful) == 3, ( - f"Expected all 3 to succeed, got {len(successful)} successes: {results}" - ) + assert ( + len(successful) == 3 + ), f"Expected all 3 to succeed, got {len(successful)} successes: {results}" assert final_balance == 90, f"Expected balance 90, got {final_balance}" # Check transaction timestamps to confirm database-level serialization @@ -583,38 +579,38 @@ async def test_concurrent_multiple_spends_sufficient_balance(server: SpinTestSer # Verify all balances are valid intermediate states for balance in actual_balances: - assert balance in expected_possible_balances, ( - f"Invalid balance {balance}, expected one of {expected_possible_balances}" - ) + assert ( + balance in expected_possible_balances + ), f"Invalid balance {balance}, expected one of {expected_possible_balances}" # Final balance should always be 90 (150 - 60) - assert min(actual_balances) == 90, ( - f"Final balance should be 90, got {min(actual_balances)}" - ) + assert ( + min(actual_balances) == 90 + ), f"Final balance should be 90, got {min(actual_balances)}" # The final transaction should always have balance 90 # The other transactions should have valid intermediate balances - assert 90 in actual_balances, ( - f"Final balance 90 should be in actual_balances: {actual_balances}" - ) + assert ( + 90 in actual_balances + ), f"Final balance 90 should be in actual_balances: {actual_balances}" # All balances should be >= 90 (the final state) - assert all(balance >= 90 for balance in actual_balances), ( - f"All balances should be >= 90, got {actual_balances}" - ) + assert all( + balance >= 90 for balance in actual_balances + ), f"All balances should be >= 90, got {actual_balances}" # CRITICAL: Transactions are atomic but can complete in any order # What matters is that all running balances are valid intermediate states # Each balance should be between 90 (final) and 140 (after first transaction) for balance in actual_balances: - assert 90 <= balance <= 140, ( - f"Balance {balance} is outside valid range [90, 140]" - ) + assert ( + 90 <= balance <= 140 + ), f"Balance {balance} is outside valid range [90, 140]" # Final balance (minimum) should always be 90 - assert min(actual_balances) == 90, ( - f"Final balance should be 90, got {min(actual_balances)}" - ) + assert ( + min(actual_balances) == 90 + ), f"Final balance should be 90, got {min(actual_balances)}" finally: await cleanup_test_user(user_id) @@ -730,9 +726,9 @@ async def test_prove_database_locking_behavior(server: SpinTestServer): print(f"\nšŸ’° Final balance: {final_balance}") if len(successful) == 3: - assert final_balance == 0, ( - f"If all succeeded, balance should be 0, got {final_balance}" - ) + assert ( + final_balance == 0 + ), f"If all succeeded, balance should be 0, got {final_balance}" print( "āœ… CONCLUSION: Database row locking causes requests to WAIT and execute serially" ) diff --git a/autogpt_platform/backend/backend/data/credit_refund_test.py b/autogpt_platform/backend/backend/data/credit_refund_test.py index 247b6e1e7a..a2031675cd 100644 --- a/autogpt_platform/backend/backend/data/credit_refund_test.py +++ b/autogpt_platform/backend/backend/data/credit_refund_test.py @@ -115,9 +115,9 @@ async def test_deduct_credits_atomic(server: SpinTestServer): where={"userId": REFUND_TEST_USER_ID} ) assert user_balance is not None - assert user_balance.balance == 500, ( - f"Expected balance 500, got {user_balance.balance}" - ) + assert ( + user_balance.balance == 500 + ), f"Expected balance 500, got {user_balance.balance}" # Verify refund transaction was created refund_tx = await CreditTransaction.prisma().find_first( @@ -211,9 +211,9 @@ async def test_handle_dispute_with_sufficient_balance( where={"userId": REFUND_TEST_USER_ID} ) assert user_balance is not None - assert user_balance.balance == 1000, ( - f"Balance should remain 1000, got {user_balance.balance}" - ) + assert ( + user_balance.balance == 1000 + ), f"Balance should remain 1000, got {user_balance.balance}" finally: await cleanup_test_user() @@ -338,9 +338,9 @@ async def test_concurrent_refunds(server: SpinTestServer): print(f"DEBUG: Final balance = {user_balance.balance}, expected = 500") # With atomic implementation, all 5 refunds should process correctly - assert user_balance.balance == 500, ( - f"Expected balance 500 after 5 refunds of 100 each, got {user_balance.balance}" - ) + assert ( + user_balance.balance == 500 + ), f"Expected balance 500 after 5 refunds of 100 each, got {user_balance.balance}" # Verify all refund transactions exist refund_txs = await CreditTransaction.prisma().find_many( @@ -349,9 +349,9 @@ async def test_concurrent_refunds(server: SpinTestServer): "type": CreditTransactionType.REFUND, } ) - assert len(refund_txs) == 5, ( - f"Expected 5 refund transactions, got {len(refund_txs)}" - ) + assert ( + len(refund_txs) == 5 + ), f"Expected 5 refund transactions, got {len(refund_txs)}" running_balances: set[int] = { tx.runningBalance for tx in refund_txs if tx.runningBalance is not None @@ -359,20 +359,20 @@ async def test_concurrent_refunds(server: SpinTestServer): # Verify all balances are valid intermediate states for balance in running_balances: - assert 500 <= balance <= 1000, ( - f"Invalid balance {balance}, should be between 500 and 1000" - ) + assert ( + 500 <= balance <= 1000 + ), f"Invalid balance {balance}, should be between 500 and 1000" # Final balance should be present - assert 500 in running_balances, ( - f"Final balance 500 should be in {running_balances}" - ) + assert ( + 500 in running_balances + ), f"Final balance 500 should be in {running_balances}" # All balances should be unique and form a valid sequence sorted_balances = sorted(running_balances, reverse=True) - assert len(sorted_balances) == 5, ( - f"Expected 5 unique balances, got {len(sorted_balances)}" - ) + assert ( + len(sorted_balances) == 5 + ), f"Expected 5 unique balances, got {len(sorted_balances)}" finally: await cleanup_test_user() diff --git a/autogpt_platform/backend/backend/data/credit_underflow_test.py b/autogpt_platform/backend/backend/data/credit_underflow_test.py index 3eaef71229..001b8045e6 100644 --- a/autogpt_platform/backend/backend/data/credit_underflow_test.py +++ b/autogpt_platform/backend/backend/data/credit_underflow_test.py @@ -12,11 +12,7 @@ import pytest from prisma.enums import CreditTransactionType from prisma.errors import UniqueViolationError from prisma.models import CreditTransaction, User, UserBalance -from prisma.types import ( - UserBalanceCreateInput, - UserBalanceUpsertInput, - UserCreateInput, -) +from prisma.types import UserBalanceCreateInput, UserBalanceUpsertInput, UserCreateInput from backend.data.credit import POSTGRES_INT_MIN, UserCredit from backend.util.test import SpinTestServer @@ -90,7 +86,9 @@ async def test_debug_underflow_step_by_step(server: SpinTestServer): # Test 2: Apply amount that should cause underflow print("\n=== Test 2: Testing underflow protection ===") - test_amount = -200 # This should cause underflow: (POSTGRES_INT_MIN + 100) + (-200) = POSTGRES_INT_MIN - 100 + test_amount = ( + -200 + ) # This should cause underflow: (POSTGRES_INT_MIN + 100) + (-200) = POSTGRES_INT_MIN - 100 expected_without_protection = current_balance + test_amount print(f"Current balance: {current_balance}") print(f"Test amount: {test_amount}") @@ -107,9 +105,9 @@ async def test_debug_underflow_step_by_step(server: SpinTestServer): print(f"Actual result: {balance_result}") # Check if underflow protection worked - assert balance_result == POSTGRES_INT_MIN, ( - f"Expected underflow protection to clamp balance to {POSTGRES_INT_MIN}, got {balance_result}" - ) + assert ( + balance_result == POSTGRES_INT_MIN + ), f"Expected underflow protection to clamp balance to {POSTGRES_INT_MIN}, got {balance_result}" # Test 3: Edge case - exactly at POSTGRES_INT_MIN print("\n=== Test 3: Testing exact POSTGRES_INT_MIN boundary ===") @@ -134,9 +132,9 @@ async def test_debug_underflow_step_by_step(server: SpinTestServer): ) print(f"After subtracting 1: {edge_result}") - assert edge_result == POSTGRES_INT_MIN, ( - f"Expected balance to remain clamped at {POSTGRES_INT_MIN}, got {edge_result}" - ) + assert ( + edge_result == POSTGRES_INT_MIN + ), f"Expected balance to remain clamped at {POSTGRES_INT_MIN}, got {edge_result}" finally: await cleanup_test_user(user_id) @@ -180,18 +178,18 @@ async def test_underflow_protection_large_refunds(server: SpinTestServer): ) # Balance should be clamped to POSTGRES_INT_MIN, not the calculated underflow value - assert final_balance == POSTGRES_INT_MIN, ( - f"Balance should be clamped to {POSTGRES_INT_MIN}, got {final_balance}" - ) - assert final_balance > expected_without_protection, ( - f"Balance should be greater than underflow result {expected_without_protection}, got {final_balance}" - ) + assert ( + final_balance == POSTGRES_INT_MIN + ), f"Balance should be clamped to {POSTGRES_INT_MIN}, got {final_balance}" + assert ( + final_balance > expected_without_protection + ), f"Balance should be greater than underflow result {expected_without_protection}, got {final_balance}" # Verify with get_credits too stored_balance = await credit_system.get_credits(user_id) - assert stored_balance == POSTGRES_INT_MIN, ( - f"Stored balance should be {POSTGRES_INT_MIN}, got {stored_balance}" - ) + assert ( + stored_balance == POSTGRES_INT_MIN + ), f"Stored balance should be {POSTGRES_INT_MIN}, got {stored_balance}" # Verify transaction was created with the underflow-protected balance transactions = await CreditTransaction.prisma().find_many( @@ -199,9 +197,9 @@ async def test_underflow_protection_large_refunds(server: SpinTestServer): order={"createdAt": "desc"}, ) assert len(transactions) > 0, "Refund transaction should be created" - assert transactions[0].runningBalance == POSTGRES_INT_MIN, ( - f"Transaction should show clamped balance {POSTGRES_INT_MIN}, got {transactions[0].runningBalance}" - ) + assert ( + transactions[0].runningBalance == POSTGRES_INT_MIN + ), f"Transaction should show clamped balance {POSTGRES_INT_MIN}, got {transactions[0].runningBalance}" finally: await cleanup_test_user(user_id) @@ -240,12 +238,12 @@ async def test_multiple_large_refunds_cumulative_underflow(server: SpinTestServe expected_balance_1 = ( initial_balance + refund_amount ) # Should be POSTGRES_INT_MIN + 200 - assert balance_1 == expected_balance_1, ( - f"First refund should result in {expected_balance_1}, got {balance_1}" - ) - assert balance_1 >= POSTGRES_INT_MIN, ( - f"First refund should not go below {POSTGRES_INT_MIN}, got {balance_1}" - ) + assert ( + balance_1 == expected_balance_1 + ), f"First refund should result in {expected_balance_1}, got {balance_1}" + assert ( + balance_1 >= POSTGRES_INT_MIN + ), f"First refund should not go below {POSTGRES_INT_MIN}, got {balance_1}" # Second refund: (POSTGRES_INT_MIN + 200) + (-300) = POSTGRES_INT_MIN - 100 (would underflow) balance_2, _ = await credit_system._add_transaction( @@ -256,9 +254,9 @@ async def test_multiple_large_refunds_cumulative_underflow(server: SpinTestServe ) # Should be clamped to minimum due to underflow protection - assert balance_2 == POSTGRES_INT_MIN, ( - f"Second refund should be clamped to {POSTGRES_INT_MIN}, got {balance_2}" - ) + assert ( + balance_2 == POSTGRES_INT_MIN + ), f"Second refund should be clamped to {POSTGRES_INT_MIN}, got {balance_2}" # Third refund: Should stay at minimum balance_3, _ = await credit_system._add_transaction( @@ -269,15 +267,15 @@ async def test_multiple_large_refunds_cumulative_underflow(server: SpinTestServe ) # Should still be at minimum - assert balance_3 == POSTGRES_INT_MIN, ( - f"Third refund should stay at {POSTGRES_INT_MIN}, got {balance_3}" - ) + assert ( + balance_3 == POSTGRES_INT_MIN + ), f"Third refund should stay at {POSTGRES_INT_MIN}, got {balance_3}" # Final balance check final_balance = await credit_system.get_credits(user_id) - assert final_balance == POSTGRES_INT_MIN, ( - f"Final balance should be {POSTGRES_INT_MIN}, got {final_balance}" - ) + assert ( + final_balance == POSTGRES_INT_MIN + ), f"Final balance should be {POSTGRES_INT_MIN}, got {final_balance}" finally: await cleanup_test_user(user_id) @@ -329,35 +327,35 @@ async def test_concurrent_large_refunds_no_underflow(server: SpinTestServer): for i, result in enumerate(results): if isinstance(result, tuple): balance, _ = result - assert balance >= POSTGRES_INT_MIN, ( - f"Result {i} balance {balance} underflowed below {POSTGRES_INT_MIN}" - ) + assert ( + balance >= POSTGRES_INT_MIN + ), f"Result {i} balance {balance} underflowed below {POSTGRES_INT_MIN}" valid_results.append(balance) elif isinstance(result, str) and "FAILED" in result: # Some operations might fail due to validation, that's okay pass else: # Unexpected exception - assert not isinstance(result, Exception), ( - f"Unexpected exception in result {i}: {result}" - ) + assert not isinstance( + result, Exception + ), f"Unexpected exception in result {i}: {result}" # At least one operation should succeed - assert len(valid_results) > 0, ( - f"At least one refund should succeed, got results: {results}" - ) + assert ( + len(valid_results) > 0 + ), f"At least one refund should succeed, got results: {results}" # All successful results should be >= POSTGRES_INT_MIN for balance in valid_results: - assert balance >= POSTGRES_INT_MIN, ( - f"Balance {balance} should not be below {POSTGRES_INT_MIN}" - ) + assert ( + balance >= POSTGRES_INT_MIN + ), f"Balance {balance} should not be below {POSTGRES_INT_MIN}" # Final balance should be valid and at or above POSTGRES_INT_MIN final_balance = await credit_system.get_credits(user_id) - assert final_balance >= POSTGRES_INT_MIN, ( - f"Final balance {final_balance} should not underflow below {POSTGRES_INT_MIN}" - ) + assert ( + final_balance >= POSTGRES_INT_MIN + ), f"Final balance {final_balance} should not underflow below {POSTGRES_INT_MIN}" finally: await cleanup_test_user(user_id) diff --git a/autogpt_platform/backend/backend/data/credit_user_balance_migration_test.py b/autogpt_platform/backend/backend/data/credit_user_balance_migration_test.py index 072fd708ed..4cf27ca022 100644 --- a/autogpt_platform/backend/backend/data/credit_user_balance_migration_test.py +++ b/autogpt_platform/backend/backend/data/credit_user_balance_migration_test.py @@ -61,9 +61,9 @@ async def test_user_balance_migration_complete(server: SpinTestServer): # User.balance should not exist or should be None/0 if it exists user_balance_attr = getattr(user, "balance", None) if user_balance_attr is not None: - assert user_balance_attr == 0 or user_balance_attr is None, ( - f"User.balance should be 0 or None, got {user_balance_attr}" - ) + assert ( + user_balance_attr == 0 or user_balance_attr is None + ), f"User.balance should be 0 or None, got {user_balance_attr}" # 2. Perform various credit operations using internal method (bypasses Stripe) await credit_system._add_transaction( @@ -88,9 +88,9 @@ async def test_user_balance_migration_complete(server: SpinTestServer): # 3. Verify UserBalance table has correct values user_balance = await UserBalance.prisma().find_unique(where={"userId": user_id}) assert user_balance is not None - assert user_balance.balance == 700, ( - f"UserBalance should be 700, got {user_balance.balance}" - ) + assert ( + user_balance.balance == 700 + ), f"UserBalance should be 700, got {user_balance.balance}" # 4. CRITICAL: Verify User.balance is NEVER updated during operations user_after = await User.prisma().find_unique(where={"id": user_id}) @@ -98,15 +98,15 @@ async def test_user_balance_migration_complete(server: SpinTestServer): user_balance_after = getattr(user_after, "balance", None) if user_balance_after is not None: # If User.balance exists, it should still be 0 (never updated) - assert user_balance_after == 0 or user_balance_after is None, ( - f"User.balance should remain 0/None after operations, got {user_balance_after}. This indicates User.balance is still being used!" - ) + assert ( + user_balance_after == 0 or user_balance_after is None + ), f"User.balance should remain 0/None after operations, got {user_balance_after}. This indicates User.balance is still being used!" # 5. Verify get_credits always returns UserBalance value, not User.balance final_balance = await credit_system.get_credits(user_id) - assert final_balance == user_balance.balance, ( - f"get_credits should return UserBalance value {user_balance.balance}, got {final_balance}" - ) + assert ( + final_balance == user_balance.balance + ), f"get_credits should return UserBalance value {user_balance.balance}, got {final_balance}" finally: await cleanup_test_user(user_id) @@ -127,9 +127,9 @@ async def test_detect_stale_user_balance_queries(server: SpinTestServer): # Verify that get_credits returns UserBalance value (5000), not any stale User.balance value balance = await credit_system.get_credits(user_id) - assert balance == 5000, ( - f"Expected get_credits to return 5000 from UserBalance, got {balance}" - ) + assert ( + balance == 5000 + ), f"Expected get_credits to return 5000 from UserBalance, got {balance}" # Verify all operations use UserBalance using internal method (bypasses Stripe) await credit_system._add_transaction( @@ -144,9 +144,9 @@ async def test_detect_stale_user_balance_queries(server: SpinTestServer): # Verify UserBalance table has the correct value user_balance = await UserBalance.prisma().find_unique(where={"userId": user_id}) assert user_balance is not None - assert user_balance.balance == 6000, ( - f"UserBalance should be 6000, got {user_balance.balance}" - ) + assert ( + user_balance.balance == 6000 + ), f"UserBalance should be 6000, got {user_balance.balance}" finally: await cleanup_test_user(user_id) @@ -199,9 +199,9 @@ async def test_concurrent_operations_use_userbalance_only(server: SpinTestServer # Verify UserBalance has correct value user_balance = await UserBalance.prisma().find_unique(where={"userId": user_id}) assert user_balance is not None - assert user_balance.balance == 400, ( - f"UserBalance should be 400, got {user_balance.balance}" - ) + assert ( + user_balance.balance == 400 + ), f"UserBalance should be 400, got {user_balance.balance}" # Critical: If User.balance exists and was used, it might have wrong value try: diff --git a/autogpt_platform/backend/backend/data/execution.py b/autogpt_platform/backend/backend/data/execution.py index 7431d4103c..b55862f6ab 100644 --- a/autogpt_platform/backend/backend/data/execution.py +++ b/autogpt_platform/backend/backend/data/execution.py @@ -36,7 +36,6 @@ from prisma.types import ( AgentNodeExecutionKeyValueDataCreateInput, AgentNodeExecutionUpdateInput, AgentNodeExecutionWhereInput, - AgentNodeExecutionWhereUniqueInput, ) from pydantic import BaseModel, ConfigDict, JsonValue, ValidationError from pydantic.fields import Field diff --git a/autogpt_platform/backend/backend/data/onboarding.py b/autogpt_platform/backend/backend/data/onboarding.py index 51fd5cf611..3cbea503b5 100644 --- a/autogpt_platform/backend/backend/data/onboarding.py +++ b/autogpt_platform/backend/backend/data/onboarding.py @@ -114,7 +114,22 @@ async def update_user_onboarding(user_id: str, data: UserOnboardingUpdate): if data.onboardingAgentExecutionId is not None: update["onboardingAgentExecutionId"] = data.onboardingAgentExecutionId - create_input = UserOnboardingCreateInput(userId=user_id, **update) + # Build create_input manually to avoid type issues with Prisma update types + create_input = UserOnboardingCreateInput( + userId=user_id, + walletShown=data.walletShown if data.walletShown else None, + notified=( + list(set(data.notified + onboarding.notified)) + if data.notified is not None + else None + ), + usageReason=data.usageReason, + integrations=data.integrations, + otherIntegrations=data.otherIntegrations, + selectedStoreListingVersionId=data.selectedStoreListingVersionId, + agentInput=SafeJson(data.agentInput) if data.agentInput is not None else None, + onboardingAgentExecutionId=data.onboardingAgentExecutionId, + ) return await UserOnboarding.prisma().upsert( where={"userId": user_id}, data=UserOnboardingUpsertInput( diff --git a/autogpt_platform/backend/test/e2e_test_data.py b/autogpt_platform/backend/test/e2e_test_data.py index d7576cdad3..cd067abbdf 100644 --- a/autogpt_platform/backend/test/e2e_test_data.py +++ b/autogpt_platform/backend/test/e2e_test_data.py @@ -22,6 +22,7 @@ import random from typing import Any, Dict, List from faker import Faker +from prisma.types import AgentBlockCreateInput # Import API functions from the backend from backend.api.features.library.db import create_library_agent, create_preset @@ -179,12 +180,12 @@ class TestDataCreator: for block in blocks_to_create: try: await prisma.agentblock.create( - data={ - "id": block.id, - "name": block.name, - "inputSchema": "{}", - "outputSchema": "{}", - } + data=AgentBlockCreateInput( + id=block.id, + name=block.name, + inputSchema="{}", + outputSchema="{}", + ) ) except Exception as e: print(f"Error creating block {block.name}: {e}") diff --git a/autogpt_platform/backend/test/test_data_creator.py b/autogpt_platform/backend/test/test_data_creator.py index befb1dcacd..82ba1f220d 100644 --- a/autogpt_platform/backend/test/test_data_creator.py +++ b/autogpt_platform/backend/test/test_data_creator.py @@ -30,13 +30,19 @@ from prisma.types import ( AgentGraphCreateInput, AgentNodeCreateInput, AgentNodeLinkCreateInput, + AgentPresetCreateInput, AnalyticsDetailsCreateInput, AnalyticsMetricsCreateInput, + APIKeyCreateInput, CreditTransactionCreateInput, IntegrationWebhookCreateInput, + LibraryAgentCreateInput, ProfileCreateInput, + StoreListingCreateInput, StoreListingReviewCreateInput, + StoreListingVersionCreateInput, UserCreateInput, + UserOnboardingCreateInput, ) faker = Faker() @@ -172,14 +178,14 @@ async def main(): for _ in range(num_presets): # Create 1 AgentPreset per user graph = random.choice(agent_graphs) preset = await db.agentpreset.create( - data={ - "name": faker.sentence(nb_words=3), - "description": faker.text(max_nb_chars=200), - "userId": user.id, - "agentGraphId": graph.id, - "agentGraphVersion": graph.version, - "isActive": True, - } + data=AgentPresetCreateInput( + name=faker.sentence(nb_words=3), + description=faker.text(max_nb_chars=200), + userId=user.id, + agentGraphId=graph.id, + agentGraphVersion=graph.version, + isActive=True, + ) ) agent_presets.append(preset) @@ -220,18 +226,18 @@ async def main(): ) library_agent = await db.libraryagent.create( - data={ - "userId": user.id, - "agentGraphId": graph.id, - "agentGraphVersion": graph.version, - "creatorId": creator_profile.id if creator_profile else None, - "imageUrl": get_image() if random.random() < 0.5 else None, - "useGraphIsActiveVersion": random.choice([True, False]), - "isFavorite": random.choice([True, False]), - "isCreatedByUser": random.choice([True, False]), - "isArchived": random.choice([True, False]), - "isDeleted": random.choice([True, False]), - } + data=LibraryAgentCreateInput( + userId=user.id, + agentGraphId=graph.id, + agentGraphVersion=graph.version, + creatorId=creator_profile.id if creator_profile else None, + imageUrl=get_image() if random.random() < 0.5 else None, + useGraphIsActiveVersion=random.choice([True, False]), + isFavorite=random.choice([True, False]), + isCreatedByUser=random.choice([True, False]), + isArchived=random.choice([True, False]), + isDeleted=random.choice([True, False]), + ) ) library_agents.append(library_agent) @@ -392,13 +398,13 @@ async def main(): user = random.choice(users) slug = faker.slug() listing = await db.storelisting.create( - data={ - "agentGraphId": graph.id, - "agentGraphVersion": graph.version, - "owningUserId": user.id, - "hasApprovedVersion": random.choice([True, False]), - "slug": slug, - } + data=StoreListingCreateInput( + agentGraphId=graph.id, + agentGraphVersion=graph.version, + owningUserId=user.id, + hasApprovedVersion=random.choice([True, False]), + slug=slug, + ) ) store_listings.append(listing) @@ -408,26 +414,26 @@ async def main(): for listing in store_listings: graph = [g for g in agent_graphs if g.id == listing.agentGraphId][0] version = await db.storelistingversion.create( - data={ - "agentGraphId": graph.id, - "agentGraphVersion": graph.version, - "name": graph.name or faker.sentence(nb_words=3), - "subHeading": faker.sentence(), - "videoUrl": get_video_url() if random.random() < 0.3 else None, - "imageUrls": [get_image() for _ in range(3)], - "description": faker.text(), - "categories": [faker.word() for _ in range(3)], - "isFeatured": random.choice([True, False]), - "isAvailable": True, - "storeListingId": listing.id, - "submissionStatus": random.choice( + data=StoreListingVersionCreateInput( + agentGraphId=graph.id, + agentGraphVersion=graph.version, + name=graph.name or faker.sentence(nb_words=3), + subHeading=faker.sentence(), + videoUrl=get_video_url() if random.random() < 0.3 else None, + imageUrls=[get_image() for _ in range(3)], + description=faker.text(), + categories=[faker.word() for _ in range(3)], + isFeatured=random.choice([True, False]), + isAvailable=True, + storeListingId=listing.id, + submissionStatus=random.choice( [ prisma.enums.SubmissionStatus.PENDING, prisma.enums.SubmissionStatus.APPROVED, prisma.enums.SubmissionStatus.REJECTED, ] ), - } + ) ) store_listing_versions.append(version) @@ -469,51 +475,49 @@ async def main(): try: await db.useronboarding.create( - data={ - "userId": user.id, - "completedSteps": completed_steps, - "walletShown": random.choice([True, False]), - "notified": ( + data=UserOnboardingCreateInput( + userId=user.id, + completedSteps=completed_steps, + walletShown=random.choice([True, False]), + notified=( random.sample(completed_steps, k=min(3, len(completed_steps))) if completed_steps else [] ), - "rewardedFor": ( + rewardedFor=( random.sample(completed_steps, k=min(2, len(completed_steps))) if completed_steps else [] ), - "usageReason": ( + usageReason=( random.choice(["personal", "business", "research", "learning"]) if random.random() < 0.7 else None ), - "integrations": random.sample( + integrations=random.sample( ["github", "google", "discord", "slack"], k=random.randint(0, 2) ), - "otherIntegrations": ( - faker.word() if random.random() < 0.2 else None - ), - "selectedStoreListingVersionId": ( + otherIntegrations=(faker.word() if random.random() < 0.2 else None), + selectedStoreListingVersionId=( random.choice(store_listing_versions).id if store_listing_versions and random.random() < 0.5 else None ), - "onboardingAgentExecutionId": ( + onboardingAgentExecutionId=( random.choice(agent_graph_executions).id if agent_graph_executions and random.random() < 0.3 else None ), - "agentRuns": random.randint(0, 10), - } + agentRuns=random.randint(0, 10), + ) ) except Exception as e: print(f"Error creating onboarding for user {user.id}: {e}") # Try simpler version await db.useronboarding.create( - data={ - "userId": user.id, - } + data=UserOnboardingCreateInput( + userId=user.id, + ) ) # Insert IntegrationWebhooks for some users @@ -544,20 +548,20 @@ async def main(): for user in users: api_key = APIKeySmith().generate_key() await db.apikey.create( - data={ - "name": faker.word(), - "head": api_key.head, - "tail": api_key.tail, - "hash": api_key.hash, - "salt": api_key.salt, - "status": prisma.enums.APIKeyStatus.ACTIVE, - "permissions": [ + data=APIKeyCreateInput( + name=faker.word(), + head=api_key.head, + tail=api_key.tail, + hash=api_key.hash, + salt=api_key.salt, + status=prisma.enums.APIKeyStatus.ACTIVE, + permissions=[ prisma.enums.APIKeyPermission.EXECUTE_GRAPH, prisma.enums.APIKeyPermission.READ_GRAPH, ], - "description": faker.text(), - "userId": user.id, - } + description=faker.text(), + userId=user.id, + ) ) # Refresh materialized views diff --git a/autogpt_platform/backend/test/test_data_updater.py b/autogpt_platform/backend/test/test_data_updater.py index 561a5da21a..b89b3478cf 100755 --- a/autogpt_platform/backend/test/test_data_updater.py +++ b/autogpt_platform/backend/test/test_data_updater.py @@ -16,6 +16,7 @@ from datetime import datetime, timedelta import prisma.enums from faker import Faker from prisma import Json, Prisma +from prisma.types import CreditTransactionCreateInput, StoreListingReviewCreateInput faker = Faker() @@ -166,16 +167,16 @@ async def main(): score = random.choices([1, 2, 3, 4, 5], weights=[5, 10, 20, 40, 25])[0] await db.storelistingreview.create( - data={ - "storeListingVersionId": version.id, - "reviewByUserId": reviewer.id, - "score": score, - "comments": ( + data=StoreListingReviewCreateInput( + storeListingVersionId=version.id, + reviewByUserId=reviewer.id, + score=score, + comments=( faker.text(max_nb_chars=200) if random.random() < 0.7 else None ), - } + ) ) new_reviews_count += 1 @@ -244,17 +245,17 @@ async def main(): ) await db.credittransaction.create( - data={ - "userId": user.id, - "amount": amount, - "type": transaction_type, - "metadata": Json( + data=CreditTransactionCreateInput( + userId=user.id, + amount=amount, + type=transaction_type, + metadata=Json( { "source": "test_updater", "timestamp": datetime.now().isoformat(), } ), - } + ) ) transaction_count += 1