"""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()