igny8/backend/test_payment_workflow.py

#!/usr/bin/env python3
"""
End-to-End Payment Workflow Test Script
Tests the complete manual payment approval flow
"""

import os
import sys
import django

# Setup Django
sys.path.insert(0, os.path.dirname(__file__))
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'igny8_core.settings')
django.setup()

from django.contrib.auth import get_user_model
from django.db import transaction
from django.utils import timezone
from decimal import Decimal
from datetime import timedelta

from igny8_core.auth.models import Account, Subscription, Plan
from igny8_core.business.billing.models import (
    Invoice, Payment, AccountPaymentMethod, CreditTransaction
)
from igny8_core.business.billing.services.invoice_service import InvoiceService

User = get_user_model()

class Colors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'

def print_header(text):
    print(f"\n{Colors.HEADER}{Colors.BOLD}{'='*60}{Colors.ENDC}")
    print(f"{Colors.HEADER}{Colors.BOLD}{text:^60}{Colors.ENDC}")
    print(f"{Colors.HEADER}{Colors.BOLD}{'='*60}{Colors.ENDC}\n")

def print_success(text):
    print(f"{Colors.OKGREEN}✓ {text}{Colors.ENDC}")

def print_error(text):
    print(f"{Colors.FAIL}✗ {text}{Colors.ENDC}")

def print_info(text):
    print(f"{Colors.OKCYAN}→ {text}{Colors.ENDC}")

def cleanup_test_data():
    """Remove test data from previous runs"""
    print_header("CLEANUP TEST DATA")
    
    # Delete test accounts
    test_emails = [
        'workflow_test_free@example.com',
        'workflow_test_paid@example.com'
    ]
    
    for email in test_emails:
        try:
            user = User.objects.filter(email=email).first()
            if user:
                # Delete associated account (cascade will handle related objects)
                account = Account.objects.filter(owner=user).first()
                if account:
                    account.delete()
                    print_success(f"Deleted account for {email}")
                user.delete()
                print_success(f"Deleted user {email}")
        except Exception as e:
            print_error(f"Error cleaning up {email}: {e}")

def test_free_trial_signup():
    """Test free trial user registration"""
    print_header("TEST 1: FREE TRIAL SIGNUP")
    
    try:
        # Get free plan
        free_plan = Plan.objects.get(slug='free')
        print_info(f"Free Plan: {free_plan.name} - {free_plan.included_credits} credits")
        
        # Create user
        with transaction.atomic():
            user = User.objects.create_user(
                username='workflow_test_free',
                email='workflow_test_free@example.com',
                password='TestPass123!',
                first_name='Free',
                last_name='Trial'
            )
            print_success(f"Created user: {user.email}")
            
            # Create account
            account = Account.objects.create(
                name=f"{user.first_name}'s Account",
                slug=f'free-trial-{timezone.now().timestamp()}',
                owner=user,
                plan=free_plan,
                status='trial',
                credits=free_plan.included_credits
            )
            print_success(f"Created account: {account.name} (ID: {account.id})")
            
            # Create credit transaction
            CreditTransaction.objects.create(
                account=account,
                transaction_type='plan_allocation',
                amount=free_plan.included_credits,
                balance_after=account.credits,
                description=f'Initial credits from {free_plan.name} plan'
            )
            print_success(f"Allocated {free_plan.included_credits} credits")
        
        # Verify
        account.refresh_from_db()
        assert account.status == 'trial', "Status should be 'trial'"
        assert account.credits == 1000, "Credits should be 1000"
        assert account.plan.slug == 'free', "Plan should be 'free'"
        
        # Check no subscription or invoice created
        sub_count = Subscription.objects.filter(account=account).count()
        invoice_count = Invoice.objects.filter(account=account).count()
        
        assert sub_count == 0, "Free trial should not have subscription"
        assert invoice_count == 0, "Free trial should not have invoice"
        
        print_success("No subscription created (correct for free trial)")
        print_success("No invoice created (correct for free trial)")
        
        print_success("\nFREE TRIAL TEST PASSED ✓")
        return account
        
    except Exception as e:
        print_error(f"Free trial test failed: {e}")
        raise

def test_paid_signup():
    """Test paid user registration with manual payment"""
    print_header("TEST 2: PAID SIGNUP WORKFLOW")
    
    try:
        # Get starter plan
        starter_plan = Plan.objects.get(slug='starter')
        print_info(f"Starter Plan: {starter_plan.name} - ${starter_plan.price} - {starter_plan.included_credits} credits")
        
        # Step 1: Create user with billing info
        print_info("\nStep 1: User Registration")
        with transaction.atomic():
            user = User.objects.create_user(
                username='workflow_test_paid',
                email='workflow_test_paid@example.com',
                password='TestPass123!',
                first_name='Paid',
                last_name='User'
            )
            print_success(f"Created user: {user.email}")
            
            # Create account with billing info
            account = Account.objects.create(
                name=f"{user.first_name}'s Account",
                slug=f'paid-user-{timezone.now().timestamp()}',
                owner=user,
                plan=starter_plan,
                status='pending_payment',
                credits=0,  # No credits until payment approved
                billing_email='billing@example.com',
                billing_address_line1='123 Main Street',
                billing_city='Karachi',
                billing_country='PK'
            )
            print_success(f"Created account: {account.name} (ID: {account.id})")
            print_info(f"  Status: {account.status}")
            print_info(f"  Credits: {account.credits}")
            
            # Create subscription
            subscription = Subscription.objects.create(
                account=account,
                plan=starter_plan,
                status='pending_payment',
                current_period_start=timezone.now(),
                current_period_end=timezone.now() + timedelta(days=30)
            )
            print_success(f"Created subscription (ID: {subscription.id})")
            print_info(f"  Status: {subscription.status}")
            
            # Create invoice
            invoice_service = InvoiceService()
            invoice = invoice_service.create_subscription_invoice(
                subscription=subscription,
                billing_period_start=subscription.current_period_start,
                billing_period_end=subscription.current_period_end
            )
            print_success(f"Created invoice: {invoice.invoice_number}")
            print_info(f"  Status: {invoice.status}")
            print_info(f"  Total: ${invoice.total}")
            print_info(f"  Has billing snapshot: {'billing_snapshot' in invoice.metadata}")
            
            # Create payment method
            payment_method = AccountPaymentMethod.objects.create(
                account=account,
                type='bank_transfer',
                is_default=True
            )
            print_success(f"Created payment method: {payment_method.type}")
        
        # Step 2: User submits payment confirmation
        print_info("\nStep 2: User Payment Confirmation")
        payment = Payment.objects.create(
            invoice=invoice,
            account=account,
            amount=invoice.total,
            currency=invoice.currency,
            payment_method='bank_transfer',
            status='pending_approval',
            manual_reference='BT-TEST-20251208-001',
            manual_notes='Test payment via ABC Bank'
        )
        print_success(f"Created payment (ID: {payment.id})")
        print_info(f"  Status: {payment.status}")
        print_info(f"  Reference: {payment.manual_reference}")
        
        # Verify pending state
        account.refresh_from_db()
        subscription.refresh_from_db()
        invoice.refresh_from_db()
        
        assert account.status == 'pending_payment', "Account should be pending_payment"
        assert account.credits == 0, "Credits should be 0 before approval"
        assert subscription.status == 'pending_payment', "Subscription should be pending_payment"
        assert invoice.status == 'pending', "Invoice should be pending"
        assert payment.status == 'pending_approval', "Payment should be pending_approval"
        
        print_success("\nPending state verified ✓")
        
        # Step 3: Admin approves payment
        print_info("\nStep 3: Admin Payment Approval")
        
        # Create admin user for approval
        admin_user = User.objects.filter(is_superuser=True).first()
        if not admin_user:
            admin_user = User.objects.create_superuser(
                username='test_admin',
                email='test_admin@example.com',
                password='admin123',
                first_name='Test',
                last_name='Admin'
            )
            print_info(f"Created admin user: {admin_user.email}")
        
        with transaction.atomic():
            # Update payment
            payment.status = 'succeeded'
            payment.approved_by = admin_user
            payment.approved_at = timezone.now()
            payment.admin_notes = 'Verified in bank statement'
            payment.save()
            print_success("Payment approved")
            
            # Update invoice
            invoice.status = 'paid'
            invoice.paid_at = timezone.now()
            invoice.save()
            print_success("Invoice marked as paid")
            
            # Update subscription
            subscription.status = 'active'
            subscription.save()
            print_success("Subscription activated")
            
            # Update account and add credits
            account.status = 'active'
            account.credits = starter_plan.included_credits
            account.save()
            print_success(f"Account activated with {starter_plan.included_credits} credits")
            
            # Log credit transaction
            credit_txn = CreditTransaction.objects.create(
                account=account,
                transaction_type='plan_allocation',
                amount=starter_plan.included_credits,
                balance_after=account.credits,
                description=f'Credits from approved payment (Invoice: {invoice.invoice_number})'
            )
            print_success("Credit transaction logged")
        
        # Final verification
        print_info("\nStep 4: Final Verification")
        account.refresh_from_db()
        subscription.refresh_from_db()
        invoice.refresh_from_db()
        payment.refresh_from_db()
        
        assert account.status == 'active', "Account should be active"
        assert account.credits == 1000, "Credits should be 1000"
        assert subscription.status == 'active', "Subscription should be active"
        assert invoice.status == 'paid', "Invoice should be paid"
        assert payment.status == 'succeeded', "Payment should be succeeded"
        assert payment.approved_by == admin_user, "Payment should have approved_by"
        
        print_success(f"Account: {account.status} ✓")
        print_success(f"Credits: {account.credits} ✓")
        print_success(f"Subscription: {subscription.status} ✓")
        print_success(f"Invoice: {invoice.status} ✓")
        print_success(f"Payment: {payment.status} ✓")
        print_success(f"Approved by: {payment.approved_by.email} ✓")
        
        # Check credit transaction
        txn = CreditTransaction.objects.filter(account=account).latest('created_at')
        print_success(f"Credit Transaction: {txn.transaction_type} | {txn.amount} credits ✓")
        
        print_success("\nPAID SIGNUP TEST PASSED ✓")
        return account
        
    except Exception as e:
        print_error(f"Paid signup test failed: {e}")
        raise

def test_payment_rejection():
    """Test payment rejection flow"""
    print_header("TEST 3: PAYMENT REJECTION")
    
    try:
        # Use the paid account from previous test
        account = Account.objects.get(owner__email='workflow_test_paid@example.com')
        
        # Create a second invoice for testing rejection
        print_info("Creating second invoice for rejection test")
        
        subscription = Subscription.objects.get(account=account)
        invoice_service = InvoiceService()
        
        with transaction.atomic():
            invoice2 = invoice_service.create_subscription_invoice(
                subscription=subscription,
                billing_period_start=subscription.current_period_start + timedelta(days=30),
                billing_period_end=subscription.current_period_end + timedelta(days=30)
            )
            print_success(f"Created invoice: {invoice2.invoice_number}")
            
            # Submit payment
            payment2 = Payment.objects.create(
                invoice=invoice2,
                account=account,
                amount=invoice2.total,
                currency=invoice2.currency,
                payment_method='bank_transfer',
                status='pending_approval',
                manual_reference='BT-INVALID-REF',
                manual_notes='Test invalid payment reference'
            )
            print_success(f"Created payment (ID: {payment2.id})")
        
        # Admin rejects payment
        print_info("\nRejecting payment...")
        admin_user = User.objects.filter(is_superuser=True).first()
        
        with transaction.atomic():
            payment2.status = 'failed'
            payment2.approved_by = admin_user
            payment2.approved_at = timezone.now()
            payment2.admin_notes = 'Reference number not found in bank statement'
            payment2.save()
            print_success("Payment rejected")
        
        # Verify rejection
        payment2.refresh_from_db()
        invoice2.refresh_from_db()
        
        assert payment2.status == 'failed', "Payment should be failed"
        assert invoice2.status == 'pending', "Invoice should remain pending"
        
        print_success(f"Payment status: {payment2.status} ✓")
        print_success(f"Invoice status: {invoice2.status} ✓")
        print_success(f"Rejection reason: {payment2.admin_notes} ✓")
        
        print_success("\nPAYMENT REJECTION TEST PASSED ✓")
        
    except Exception as e:
        print_error(f"Payment rejection test failed: {e}")
        raise

def print_summary():
    """Print test summary"""
    print_header("TEST SUMMARY")
    
    # Count accounts by status
    from django.db.models import Count
    
    status_counts = Account.objects.values('status').annotate(count=Count('id'))
    print_info("Account Status Distribution:")
    for item in status_counts:
        print(f"  {item['status']:20} {item['count']} account(s)")
    
    # Count payments by status
    payment_counts = Payment.objects.values('status').annotate(count=Count('id'))
    print_info("\nPayment Status Distribution:")
    for item in payment_counts:
        print(f"  {item['status']:20} {item['count']} payment(s)")
    
    # Count subscriptions by status
    sub_counts = Subscription.objects.values('status').annotate(count=Count('id'))
    print_info("\nSubscription Status Distribution:")
    for item in sub_counts:
        print(f"  {item['status']:20} {item['count']} subscription(s)")
    
    print()

def main():
    """Run all tests"""
    print_header("PAYMENT WORKFLOW E2E TEST SUITE")
    print(f"{Colors.BOLD}Date: {timezone.now().strftime('%Y-%m-%d %H:%M:%S')}{Colors.ENDC}\n")
    
    try:
        # Cleanup
        cleanup_test_data()
        
        # Run tests
        test_free_trial_signup()
        test_paid_signup()
        test_payment_rejection()
        
        # Summary
        print_summary()
        
        # Final success
        print_header("ALL TESTS PASSED ✓")
        print(f"{Colors.OKGREEN}{Colors.BOLD}The payment workflow is functioning correctly!{Colors.ENDC}\n")
        
        return 0
        
    except Exception as e:
        print_header("TESTS FAILED ✗")
        print(f"{Colors.FAIL}{Colors.BOLD}Error: {e}{Colors.ENDC}\n")
        import traceback
        traceback.print_exc()
        return 1

if __name__ == '__main__':
    sys.exit(main())