Simple Payment API Testing

from locust import task, HttpUser
import random
import json

class PaymentTestUser(HttpUser):
    def on_start(self):
        # Payment test data
        self.test_cards = [
            {"number": "4242424242424242", "exp_month": 12, "exp_year": 2025, "cvc": "123"},
            {"number": "4000000000000002", "exp_month": 11, "exp_year": 2025, "cvc": "456"},  # Declined card
            {"number": "4000000000009995", "exp_month": 10, "exp_year": 2025, "cvc": "789"}   # Insufficient funds
        ]
        
        self.amounts = [999, 1500, 2999, 4999, 9999]  # Amounts in cents
        self.currencies = ["usd", "eur", "gbp"]
        
        # API endpoints
        self.payment_endpoints = {
            "charge": "/api/payments/charge",
            "refund": "/api/payments/refund",
            "validate": "/api/payments/validate"
        }

    @task(4)
    def test_payment_charge(self):
        """Test payment charge with valid card"""
        card = self.test_cards[0]  # Use valid card
        amount = random.choice(self.amounts)
        currency = random.choice(self.currencies)
        
        charge_data = {
            "amount": amount,
            "currency": currency,
            "card": card,
            "description": f"Test charge ${amount/100}"
        }
        
        with self.client.post(
            self.payment_endpoints["charge"],
            json=charge_data,
            name="Payment Charge"
        ) as response:
            if response.status_code == 200:
                try:
                    result = response.json()
                    if result.get("status") == "succeeded":
                        print(f"Payment successful: ${amount/100} {currency}")
                    else:
                        print(f"Payment status: {result.get('status')}")
                except json.JSONDecodeError:
                    response.failure("Invalid JSON response")
            elif response.status_code == 400:
                print("Payment failed: Bad request")
            elif response.status_code == 402:
                print("Payment failed: Card declined")
            else:
                response.failure(f"Payment charge failed: {response.status_code}")

    @task(2)
    def test_payment_declined(self):
        """Test payment with declined card"""
        declined_card = self.test_cards[1]  # Declined card
        amount = random.choice(self.amounts)
        
        charge_data = {
            "amount": amount,
            "currency": "usd",
            "card": declined_card,
            "description": "Test declined payment"
        }
        
        with self.client.post(
            self.payment_endpoints["charge"],
            json=charge_data,
            name="Payment Declined"
        ) as response:
            if response.status_code == 402:
                print("Payment correctly declined")
            elif response.status_code == 200:
                try:
                    result = response.json()
                    if result.get("status") == "failed":
                        print("Payment failed as expected")
                    else:
                        response.failure("Declined card was accepted")
                except json.JSONDecodeError:
                    response.failure("Invalid JSON response")
            else:
                print(f"Declined payment returned: {response.status_code}")

    @task(2)
    def test_payment_refund(self):
        """Test payment refund"""
        # First create a charge, then refund it
        card = self.test_cards[0]
        amount = random.choice(self.amounts)
        
        # Create charge
        charge_data = {
            "amount": amount,
            "currency": "usd",
            "card": card,
            "description": "Test charge for refund"
        }
        
        with self.client.post(
            self.payment_endpoints["charge"],
            json=charge_data,
            name="Charge for Refund"
        ) as charge_response:
            if charge_response.status_code == 200:
                try:
                    charge_result = charge_response.json()
                    charge_id = charge_result.get("id")
                    
                    if charge_id:
                        # Process refund
                        refund_data = {
                            "charge_id": charge_id,
                            "amount": amount,
                            "reason": "requested_by_customer"
                        }
                        
                        with self.client.post(
                            self.payment_endpoints["refund"],
                            json=refund_data,
                            name="Payment Refund"
                        ) as refund_response:
                            if refund_response.status_code == 200:
                                print(f"Refund successful: ${amount/100}")
                            else:
                                response.failure(f"Refund failed: {refund_response.status_code}")
                    else:
                        print("No charge ID for refund")
                        
                except json.JSONDecodeError:
                    print("Invalid JSON in charge response")

    @task(2)
    def test_card_validation(self):
        """Test payment method validation"""
        card = random.choice(self.test_cards)
        
        validation_data = {
            "card": card,
            "validate_only": True
        }
        
        with self.client.post(
            self.payment_endpoints["validate"],
            json=validation_data,
            name="Card Validation"
        ) as response:
            if response.status_code == 200:
                try:
                    result = response.json()
                    is_valid = result.get("valid", False)
                    card_type = result.get("card_type", "unknown")
                    
                    print(f"Card validation: {is_valid}, type: {card_type}")
                    
                except json.JSONDecodeError:
                    response.failure("Invalid JSON response")
            else:
                response.failure(f"Card validation failed: {response.status_code}")

    @task(1)
    def test_invalid_payment_data(self):
        """Test payment with invalid data"""
        invalid_scenarios = [
            {"amount": -100, "currency": "usd", "description": "Negative amount"},
            {"amount": 1000, "currency": "xxx", "description": "Invalid currency"},
            {"amount": 0, "currency": "usd", "description": "Zero amount"}
        ]
        
        scenario = random.choice(invalid_scenarios)
        scenario["card"] = self.test_cards[0]
        
        with self.client.post(
            self.payment_endpoints["charge"],
            json=scenario,
            name="Invalid Payment Data"
        ) as response:
            if response.status_code == 400:
                print(f"Invalid data correctly rejected: {scenario['description']}")
            elif response.status_code == 200:
                response.failure(f"Invalid data accepted: {scenario['description']}")
            else:
                print(f"Invalid data returned: {response.status_code}")