Untitled

from rest_framework.test import APITestCase
from rest_framework.reverse import reverse
from rest_framework.status import HTTP_200_OK, HTTP_400_BAD_REQUEST
from api.models import UserTradingDetails, AffiliatedUser
from extensions.auth.models import User


class AddReferralTestCase(APITestCase):

    @classmethod
    def setUpTestData(cls):
        cls.user1 = User.objects.create_user(username='testuser1', password='testpassword1')
        cls.user2 = User.objects.create_user(username='testuser2', password='testpassword2')
        cls.trading_details = UserTradingDetails.objects.create(
            user=cls.user1,
            steamid='12345678',
            api_key='api_key_123',
            trade_link='http://example.com',
            affiliate_code="VALIDCODE123"
        )

    def test_use_own_affiliate_code(self):
        self.client.force_authenticate(self.user1)
        data = {
            "affiliate_code": "VALIDCODE123"
        }
        response = self.client.post(path=reverse("add_referral"), data=data)
        self.assertEqual(HTTP_400_BAD_REQUEST, response.status_code)
        self.assertEqual(response.data, "You cannot use your own affiliate code!")

    def test_invalid_affiliate_code(self):
        self.client.force_authenticate(self.user2)
        data = {
            "affiliate_code": "INVALID_CODE"
        }
        response = self.client.post(path=reverse("add_referral"), data=data)
        self.assertEqual(HTTP_400_BAD_REQUEST, response.status_code)
        self.assertEqual(response.data, "User not found")

    def test_already_used_affiliate_code(self):
        AffiliatedUser.objects.create(user=self.user2, invited_by=self.user1)
        self.client.force_authenticate(self.user2)
        data = {
            "affiliate_code": "VALIDCODE123"
        }
        response = self.client.post(path=reverse("add_referral"), data=data)
        self.assertEqual(HTTP_400_BAD_REQUEST, response.status_code)
        self.assertEqual(response.data, "You have already used an affiliate code")

    def test_successful_use_of_affiliate_code(self):
        self.client.force_authenticate(self.user2)
        data = {
            "affiliate_code": "VALIDCODE123"
        }
        response = self.client.post(path=reverse("add_referral"), data=data)
        self.assertEqual(HTTP_200_OK, response.status_code)
        self.assertEqual(response.data, "Success")
        affiliated_user = AffiliatedUser.objects.get(user=self.user2)
        self.assertEqual(affiliated_user.invited_by, self.user1)