89 lines
2.6 KiB
Python
89 lines
2.6 KiB
Python
"""Helper utilities for domain model Robot tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from datetime import UTC, datetime
|
|
from decimal import Decimal
|
|
|
|
from cleveragents.domain.models.core.enums import ModelProvider
|
|
from cleveragents.domain.models.core.org import (
|
|
CloudBillingFields,
|
|
CreditsTransaction,
|
|
CreditsTransactionType,
|
|
CreditType,
|
|
Org,
|
|
)
|
|
|
|
|
|
def _org_test() -> None:
|
|
billing = CloudBillingFields(
|
|
credits_balance=Decimal("42.5"),
|
|
monthly_grant=Decimal("10.0"),
|
|
auto_rebuy_enabled=True,
|
|
auto_rebuy_min_threshold=Decimal("5.0"),
|
|
auto_rebuy_to_balance=Decimal("15.0"),
|
|
notify_threshold=Decimal("2.0"),
|
|
max_threshold_per_month=Decimal("100.0"),
|
|
billing_cycle_started_at=datetime(2025, 1, 1, tzinfo=UTC),
|
|
changed_billing_mode=False,
|
|
trial_paid=False,
|
|
stripe_subscription_id="sub-123",
|
|
subscription_status="active",
|
|
subscription_paused_at=None,
|
|
stripe_payment_method="pm_123",
|
|
subscription_action_required=False,
|
|
subscription_action_required_invoice_url=None,
|
|
)
|
|
org = Org(
|
|
id="org-robot",
|
|
name="Robot Org",
|
|
is_trial=True,
|
|
auto_add_domain_users=True,
|
|
integrated_models_mode=None,
|
|
cloud_billing_fields=billing,
|
|
)
|
|
assert org.cloud_billing_fields is not None
|
|
assert org.cloud_billing_fields.credits_balance == Decimal("42.5")
|
|
print("org-model-ok")
|
|
|
|
|
|
def _credits_test() -> None:
|
|
tx = CreditsTransaction(
|
|
id="txn-robot",
|
|
org_id="org-robot",
|
|
org_name="Robot Org",
|
|
user_id="user-robot",
|
|
user_email="robot@example.com",
|
|
user_name="Robot",
|
|
transaction_type=CreditsTransactionType.CREDIT,
|
|
amount=Decimal("10.0"),
|
|
start_balance=Decimal("5.0"),
|
|
end_balance=Decimal("15.0"),
|
|
credit_type=CreditType.TRIAL,
|
|
credit_is_auto_rebuy=False,
|
|
debit_model_provider=ModelProvider.OPENAI,
|
|
debit_model_name="gpt-4.1",
|
|
debit_model_role="ModelRolePlanner",
|
|
created_at=datetime(2025, 1, 2, tzinfo=UTC),
|
|
)
|
|
assert tx.debit_model_provider == ModelProvider.OPENAI
|
|
assert tx.amount == Decimal("10.0")
|
|
print("credits-model-ok")
|
|
|
|
|
|
def main() -> None:
|
|
if len(sys.argv) < 2:
|
|
raise SystemExit("Expected command argument")
|
|
command = sys.argv[1]
|
|
if command == "org":
|
|
_org_test()
|
|
elif command == "credits":
|
|
_credits_test()
|
|
else:
|
|
raise SystemExit(f"Unknown helper command: {command}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|