from datetime import timedelta from django.test import TestCase from django.utils import timezone from labs.billing.lab import redeem_lab_with_credits, refund_credits_for_lab_failure from labs.billing.voucher import redeem_voucher, VOUCHER_FREE_CREDITS_EXPIRES_IN_DAYS from labs.billing.user import INITIAL_FREE_CREDITS_AMOUNT from labs.billing.utils import round_to_end_of_day from labs.billing.exceptions import InsufficientCreditsError from labs.models import ( User, UserCreditDetail, CreditTransaction, CreditVoucher, LabInstance, ) from labs.lab_manager_v2.lab_instance import check_user_can_start_lab from labs.lab_manager_v2.lab_spec import load_lab_spec class TestBillingLab(TestCase): def setUp(self): self.user = User.objects.create_user(email="test@test.com") self.user_credit_summary = self.user.credit_summary def test_redeem_lab_with_credits_success(self): voucher_code_1 = "11111-11111-11111-11111-11111" voucher_code_2 = "22222-22222-22222-22222-22222" voucher_code_3 = "33333-33333-33333-33333-33333" tomorrow = timezone.now() + timedelta(days=1) CreditVoucher.objects.create( unique_code=voucher_code_1, credits=10, valid_before=tomorrow.date(), used=CreditVoucher.UsedChoices.N, ) CreditVoucher.objects.create( unique_code=voucher_code_2, credits=10, valid_before=tomorrow.date(), used=CreditVoucher.UsedChoices.N, ) CreditVoucher.objects.create( unique_code=voucher_code_3, credits=20, valid_before=tomorrow.date() + timedelta(days=1), used=CreditVoucher.UsedChoices.N, ) redeem_voucher(self.user, voucher_code_1) redeem_voucher(self.user, voucher_code_2) redeem_voucher(self.user, voucher_code_3) self.user_credit_summary.refresh_from_db() self.assertEqual( self.user_credit_summary.free_credits, 10 * 2 + 20 + INITIAL_FREE_CREDITS_AMOUNT, ) self.assertEqual(self.user_credit_summary.paid_credits, 0) demo_001_lab_spec = load_lab_spec("demo_001") result = check_user_can_start_lab(self.user, demo_001_lab_spec) self.assertEqual(result.can_start, False) self.assertEqual(result.access_type, LabInstance.LabAccessType.CREDITS) self.assertEqual(result.user_redemption, None) redemption = redeem_lab_with_credits( user=self.user, lab_code="demo_001", credits_to_deduct=10 ) self.assertEqual(redemption.user_id, self.user.id) self.assertEqual(redemption.lab_code, "demo_001") self.assertIsNone(redemption.lab_instance) result = check_user_can_start_lab(self.user, demo_001_lab_spec) self.assertEqual(result.can_start, True) self.assertEqual(result.access_type, LabInstance.LabAccessType.CREDITS) self.assertEqual(result.user_redemption, redemption) # Deduct soonest-expiring free credits first, then paid credits. # Voucher credits expire redeem_at+365d, which is sooner than signup (+366d), # so the earliest voucher detail (smaller id) is deducted first. voucher_1 = CreditVoucher.objects.get(unique_code=voucher_code_1) expected_details = ( UserCreditDetail.objects.filter( user=self.user, credit_type=UserCreditDetail.CreditType.FREE, initial_amount=10, ) .order_by("id") .first() ) self.assertEqual(expected_details.amount, 0) self.assertEqual(expected_details.initial_amount, 10) self.assertEqual( expected_details.expires_at, round_to_end_of_day( voucher_1.redeem_at + timedelta(days=VOUCHER_FREE_CREDITS_EXPIRES_IN_DAYS) ), ) self.assertNotEqual(expected_details.expires_at.date(), voucher_1.valid_before) self.user_credit_summary.refresh_from_db() self.assertEqual( self.user_credit_summary.free_credits, INITIAL_FREE_CREDITS_AMOUNT + 10 + 20, ) self.assertEqual(self.user_credit_summary.paid_credits, 0) transaction = ( CreditTransaction.objects.filter(user=self.user) .order_by("-created_at") .first() ) self.assertEqual(transaction.user_id, self.user.id) self.assertEqual( transaction.transaction_type, CreditTransaction.TransactionType.REDEEM_LAB ) self.assertEqual(transaction.paid_credits, 0) self.assertEqual(transaction.free_credits, -10) self.assertEqual( transaction.log["previous"], { "free_credits": INITIAL_FREE_CREDITS_AMOUNT + 10 * 2 + 20, "paid_credits": 0, }, ) self.assertEqual(transaction.redemption_id, redemption.id) def test_redeem_lab_with_credits_insufficient_credits(self): with self.assertRaises(InsufficientCreditsError): redeem_lab_with_credits( user=self.user, lab_code="demo_001", credits_to_deduct=INITIAL_FREE_CREDITS_AMOUNT + 1, ) self.user_credit_summary.refresh_from_db() self.assertEqual( self.user_credit_summary.free_credits, INITIAL_FREE_CREDITS_AMOUNT ) self.assertEqual(self.user_credit_summary.paid_credits, 0) self.assertEqual(CreditTransaction.objects.filter(user=self.user).count(), 1) self.assertEqual( CreditTransaction.objects.filter( user=self.user, transaction_type=CreditTransaction.TransactionType.REDEEM_LAB, ).count(), 0, ) def test_refund_credits_for_lab_failure(self): self.assertEqual( self.user_credit_summary.free_credits, INITIAL_FREE_CREDITS_AMOUNT ) self.assertEqual(self.user_credit_summary.paid_credits, 0) redemption = redeem_lab_with_credits( user=self.user, lab_code="demo_001", credits_to_deduct=1 ) lab_instance = LabInstance.objects.create( user=self.user, lab_code="demo_001", status=LabInstance.LabInstanceStatus.CREATING_FAILED, access_type=LabInstance.LabAccessType.CREDITS, scheduled_terminate_at=timezone.now() + timedelta(hours=2), ) redemption.lab_instance = lab_instance redemption.save() redemption_txn = CreditTransaction.objects.get( user=self.user, transaction_type=CreditTransaction.TransactionType.REDEEM_LAB, redemption=redemption, ) self.assertEqual(redemption_txn.paid_credits, 0) self.assertEqual(redemption_txn.free_credits, -1) self.assertEqual(redemption_txn.redemption_id, redemption.id) self.user_credit_summary.refresh_from_db() self.assertEqual( self.user_credit_summary.free_credits, INITIAL_FREE_CREDITS_AMOUNT - 1 ) self.assertEqual(self.user_credit_summary.paid_credits, 0) refund_credits_for_lab_failure(lab_instance=lab_instance) self.user_credit_summary.refresh_from_db() self.assertEqual( self.user_credit_summary.free_credits, INITIAL_FREE_CREDITS_AMOUNT ) self.assertEqual(self.user_credit_summary.paid_credits, 0) refund_txn = CreditTransaction.objects.get( user=self.user, transaction_type=CreditTransaction.TransactionType.REDEEM_LAB_REFUND, ) self.assertEqual(refund_txn.user_id, self.user.id) self.assertEqual(refund_txn.paid_credits, 0) self.assertEqual(refund_txn.free_credits, 1) self.assertEqual(refund_txn.reference_transaction_id, redemption_txn.id) # refund again, everything should not be changed refund_credits_for_lab_failure(lab_instance=lab_instance) self.user_credit_summary.refresh_from_db() self.assertEqual( self.user_credit_summary.free_credits, INITIAL_FREE_CREDITS_AMOUNT ) self.assertEqual(self.user_credit_summary.paid_credits, 0) self.assertEqual( CreditTransaction.objects.filter( user=self.user, transaction_type=CreditTransaction.TransactionType.REDEEM_LAB_REFUND, ).count(), 1, ) def test_refund_credits_for_lab_failure_not_redeemed_with_credits(self): not_refund_access_types = [ LabInstance.LabAccessType.INTERNAL, LabInstance.LabAccessType.WORKSHOP, ] for access_type in not_refund_access_types: with self.subTest(access_type=access_type): lab_instance = LabInstance.objects.create( user=self.user, lab_code="demo_001", status=LabInstance.LabInstanceStatus.CREATING_FAILED, access_type=access_type, scheduled_terminate_at=timezone.now() + timedelta(hours=2), ) refund_credits_for_lab_failure(lab_instance=lab_instance) self.assertEqual( CreditTransaction.objects.filter( user=self.user, transaction_type=CreditTransaction.TransactionType.REDEEM_LAB_REFUND, ).count(), 0, ) self.user_credit_summary.refresh_from_db() self.assertEqual( self.user_credit_summary.free_credits, INITIAL_FREE_CREDITS_AMOUNT ) self.assertEqual(self.user_credit_summary.paid_credits, 0) def test_refund_credits_for_lab_failure_not_creating_failed(self): lab_instance = LabInstance.objects.create( user=self.user, lab_code="demo_001", status=LabInstance.LabInstanceStatus.RUNNING, scheduled_terminate_at=timezone.now() + timedelta(hours=2), ) refund_credits_for_lab_failure(lab_instance=lab_instance) self.assertEqual( CreditTransaction.objects.filter( user=self.user, transaction_type=CreditTransaction.TransactionType.REDEEM_LAB_REFUND, ).count(), 0, ) self.user_credit_summary.refresh_from_db() self.assertEqual( self.user_credit_summary.free_credits, INITIAL_FREE_CREDITS_AMOUNT ) self.assertEqual(self.user_credit_summary.paid_credits, 0)