from datetime import timedelta from rest_framework.test import APITestCase from rest_framework import status from rest_framework.test import APIClient from django.utils import timezone from labs.models import ( User, LabInstance, LabJob, LabFeedback, UserLabRedemption, ) from labs.lab_manager_v2.lab_spec import ( load_lab_spec, LabSpec, TiDBEnvironmentType, get_all_labs, ) from labs.billing.user import INITIAL_FREE_CREDITS_AMOUNT class ActiveLabInstanceTests(APITestCase): fixtures = [ "labs/tests/fixtures/user0.json", ] ENVIRONMENT_URL = "/api/v2/lab-instances/active" FEEDBACK_URL = "/api/labs/{lab_code}/feedback" def setUp(self): self.user = User.objects.create_user(email="test@example.com") self.client = APIClient() self.client.force_login(self.user) def _create_lab_instance(self, lab_status, exited_at=None): lab_instance = LabInstance.objects.create( user=self.user, lab_code="demo_001", region_name="oregon", tidb_version="v8.1.1", exited_at=exited_at, status=lab_status, scheduled_terminate_at=timezone.now() + timedelta(hours=2), ) create_job = LabJob.objects.create( type=LabJob.TypeChoices.CREATE, status=LabJob.StatusChoices.PENDING, lab_instance=lab_instance, ) terminate_job = LabJob.objects.create( type=LabJob.TypeChoices.TERMINATE, status=LabJob.StatusChoices.PENDING, lab_instance=lab_instance, ) lab_instance.create_job = create_job lab_instance.terminate_job = terminate_job lab_instance.save() return lab_instance def test_get_lab_active_environment_unauthorized(self): self.client.logout() response = self.client.get(self.ENVIRONMENT_URL) self.assertEqual(response.status_code, status.HTTP_401_UNAUTHORIZED) def test_get_lab_active_environment_with_no_active_environment(self): response = self.client.get(self.ENVIRONMENT_URL) self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND) def test_get_lab_active_environment_with_not_exited_lab_instance(self): for lab_status in LabInstance.LabInstanceStatus: # exclude TIMEOUT status, it's a virtual status if lab_status == LabInstance.LabInstanceStatus.TIMEOUT: continue with self.subTest(lab_status=lab_status): self._create_lab_instance(lab_status) response = self.client.get(self.ENVIRONMENT_URL) self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertEqual(response.data["lab_code"], "demo_001") LabInstance.objects.all().delete() def test_get_lab_active_environment_with_exited_lab_instance(self): for lab_status in LabInstance.LabInstanceStatus: # exclude TIMEOUT status, it's a virtual status if lab_status == LabInstance.LabInstanceStatus.TIMEOUT: continue with self.subTest(lab_status=lab_status): self._create_lab_instance(lab_status, timezone.now()) response = self.client.get(self.ENVIRONMENT_URL) self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND) LabInstance.objects.all().delete() def test_get_lab_feedback_unauthorized(self): self.client.logout() response = self.client.get(self.FEEDBACK_URL.format(lab_code="demo_001")) self.assertEqual(response.status_code, status.HTTP_401_UNAUTHORIZED) def test_get_lab_feedback_not_found(self): self.assertEqual(LabFeedback.objects.count(), 0) response = self.client.get(self.FEEDBACK_URL.format(lab_code="demo_001")) self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND) def test_submit_lab_feedback_success(self): self.assertEqual(LabFeedback.objects.count(), 0) response = self.client.post( self.FEEDBACK_URL.format(lab_code="demo_001"), data={ "rating_score": 5, "comment": "This is a test comment", }, ) self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertEqual(LabFeedback.objects.count(), 1) feedback = LabFeedback.objects.first() self.assertEqual(feedback.rating_score, 5) self.assertEqual(feedback.comment, "This is a test comment") self.assertEqual(feedback.email, self.user.email) self.assertEqual(feedback.lab_code, "demo_001") response = self.client.get(self.FEEDBACK_URL.format(lab_code="demo_001")) self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertEqual(response.data["rating_score"], 5) self.assertEqual(response.data["comment"], "This is a test comment") def test_submit_lab_feedback_missing_rating(self): response = self.client.post( self.FEEDBACK_URL.format(lab_code="demo_001"), data={ "comment": "This is a test comment", }, ) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertEqual(LabFeedback.objects.count(), 0) def test_submit_lab_feedback_missing_comment(self): response = self.client.post( self.FEEDBACK_URL.format(lab_code="demo_001"), data={ "rating_score": 5, }, ) self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertEqual(LabFeedback.objects.count(), 1) feedback = LabFeedback.objects.first() self.assertEqual(feedback.rating_score, 5) self.assertEqual(feedback.comment, None) self.assertEqual(feedback.email, self.user.email) self.assertEqual(feedback.lab_code, "demo_001") def test_submit_lab_feedback_invalid_rating(self): response = self.client.post( self.FEEDBACK_URL.format(lab_code="demo_001"), data={ "rating_score": 6, }, ) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertEqual(LabFeedback.objects.count(), 0) class RedeemLabTests(APITestCase): fixtures = [ "labs/tests/fixtures/user0.json", ] REDEEM_LAB_URL = "/api/v2/labs/{lab_code}/redeem" def setUp(self): self.client = APIClient() self.user = User.objects.get(email="test1@example.com") self.pingcap_user = User.objects.get(email="employee1@pingcap.com") self.client.force_login(self.user) return super().setUp() def _get_labs(self) -> list[dict]: return get_all_labs().values() def _get_redeemable_lab(self) -> LabSpec: for lab in self._get_labs(): lab_spec = load_lab_spec(lab["lab_code"]) if lab_spec.employee_only: continue if lab_spec.credits <= 0: continue if lab_spec.credits > INITIAL_FREE_CREDITS_AMOUNT: continue if not lab_spec.credit_redeemable: continue return lab_spec else: raise ValueError("No redeemable lab found") def _get_employee_only_lab(self) -> LabSpec: for lab in self._get_labs(): lab_spec = load_lab_spec(lab["lab_code"]) if lab_spec.employee_only: return lab_spec else: raise ValueError("No employee-only lab found") def test_redeem_lab_not_logged_in(self): self.client.logout() lab = self._get_redeemable_lab() response = self.client.post(self.REDEEM_LAB_URL.format(lab_code=lab.lab_code)) self.assertEqual(response.status_code, status.HTTP_401_UNAUTHORIZED) def test_redeem_lab_with_invalid_lab_code(self): response = self.client.post( self.REDEEM_LAB_URL.format(lab_code="invalid_lab_code") ) self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND) def test_redeem_employee_only_lab(self): lab = self._get_employee_only_lab() response = self.client.post(self.REDEEM_LAB_URL.format(lab_code=lab.lab_code)) self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) def test_employee_redeem_lab(self): lab = self._get_redeemable_lab() self.client.force_login(self.pingcap_user) response = self.client.post(self.REDEEM_LAB_URL.format(lab_code=lab.lab_code)) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertIn("non_field_errors", response.data) self.assertIn( "You are a PingCAP employee, you don't need to redeem this lab with credits", str(response.data["non_field_errors"][0]), ) def test_redeem_lab_success(self): lab = self._get_redeemable_lab() response = self.client.post(self.REDEEM_LAB_URL.format(lab_code=lab.lab_code)) self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT) self.assertEqual(UserLabRedemption.objects.count(), 1) redemption = UserLabRedemption.objects.first() self.assertEqual(redemption.user, self.user) self.assertEqual(redemption.lab_code, lab.lab_code) def test_redeem_lab_insufficient_credits(self): lab = self._get_redeemable_lab() summary = self.user.credit_summary summary.free_credits = 0 summary.paid_credits = 0 summary.save() self.user.credit_details.update(amount=0) response = self.client.post(self.REDEEM_LAB_URL.format(lab_code=lab.lab_code)) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertIn("non_field_errors", response.data) self.assertIn( "You don't have enough credits to redeem this lab", str(response.data["non_field_errors"][0]), ) self.assertEqual(UserLabRedemption.objects.count(), 0) class CreateLabTests(APITestCase): fixtures = [ "labs/tests/fixtures/user0.json", ] REDEEM_LAB_URL = "/api/v2/labs/{lab_code}/redeem" CREATE_LAB_URL = "/api/v2/lab-instances" def setUp(self): self.client = APIClient() self.user = User.objects.get(email="test1@example.com") self.client.force_login(self.user) return super().setUp() def _get_labs(self) -> list[dict]: return get_all_labs().values() def _get_tidbcloud_starter_lab( self, disable_serverless_sandbox: bool = False ) -> LabSpec: for lab in self._get_labs(): lab_spec = load_lab_spec(lab["lab_code"]) if lab_spec.employee_only: continue if lab_spec.credits > INITIAL_FREE_CREDITS_AMOUNT: continue if len(lab_spec.tidb_environments) != 1: continue if lab_spec.tidb_environments[0] != TiDBEnvironmentType.TIDB_CLOUD_STARTER: continue lab_disable_serverless_sandbox = lab_spec.features.get( "disable_serverless_sandbox", False ) if lab_disable_serverless_sandbox != disable_serverless_sandbox: continue return lab_spec else: raise ValueError("No TiDB Cloud Starter lab found") def _get_self_managed_tidb_lab(self) -> LabSpec: for lab in self._get_labs(): lab_spec = load_lab_spec(lab["lab_code"]) if lab_spec.employee_only: continue if lab_spec.tidb_environments[0] == TiDBEnvironmentType.TIDB_SELF_MANAGED: return lab_spec else: raise ValueError("No Self-Managed TiDB lab found") def _create_lab_instance(self, lab_status, exited_at=None): lab_instance = LabInstance.objects.create( user=self.user, lab_code="demo_001", region_name="oregon", tidb_version="v8.1.1", exited_at=exited_at, status=lab_status, scheduled_terminate_at=timezone.now() + timedelta(hours=2), ) create_job = LabJob.objects.create( type=LabJob.TypeChoices.CREATE, status=LabJob.StatusChoices.PENDING, lab_instance=lab_instance, ) terminate_job = LabJob.objects.create( type=LabJob.TypeChoices.TERMINATE, status=LabJob.StatusChoices.PENDING, lab_instance=lab_instance, ) lab_instance.create_job = create_job lab_instance.terminate_job = terminate_job lab_instance.save() return lab_instance def test_create_lab_not_logged_in(self): self.client.logout() response = self.client.post( self.CREATE_LAB_URL, data={ "lab_code": "demo_001", "tidb_environment": "tidb_cloud_starter", "region_name": "oregon", "cloud_provider": "aws", }, ) self.assertEqual(response.status_code, status.HTTP_401_UNAUTHORIZED) def test_create_lab_with_invalid_lab_code(self): response = self.client.post( self.CREATE_LAB_URL, data={ "lab_code": "invalid_lab_code", "tidb_environment": TiDBEnvironmentType.TIDB_CLOUD_STARTER.value, "region_name": "oregon", "cloud_provider": "aws", }, ) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertIn("lab_code", response.data) def test_create_lab_with_existing_active_lab_instance(self): lab = self._get_tidbcloud_starter_lab() cloud_provider = lab.regions.get_supported_cloud_providers()[0] region_name = lab.regions.get_supported_regions(cloud_provider)[0].value self._create_lab_instance(LabInstance.LabInstanceStatus.RUNNING) response = self.client.post( self.CREATE_LAB_URL, data={ "lab_code": lab.lab_code, "tidb_environment": TiDBEnvironmentType.TIDB_CLOUD_STARTER, "cloud_provider": cloud_provider, "region_name": region_name, }, ) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertIn("non_field_errors", response.data) def test_create_lab_with_disable_serverless_sandbox_feature(self): pass def test_create_lab_with_invalid_tidb_environment(self): lab = self._get_tidbcloud_starter_lab() cloud_provider = lab.regions.get_supported_cloud_providers()[0] region_name = lab.regions.get_supported_regions(cloud_provider)[0].value response = self.client.post( self.CREATE_LAB_URL, data={ "lab_code": lab.lab_code, "tidb_environment": TiDBEnvironmentType.TIDB_CLOUD_ESSENTIAL.value, "region_name": region_name, "cloud_provider": cloud_provider, }, ) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertIn("tidb_environment", response.data) def test_create_lab_with_invalid_region_name(self): lab = self._get_tidbcloud_starter_lab() cloud_provider = lab.regions.get_supported_cloud_providers()[0] region_name = "invalid_region_name" response = self.client.post( self.CREATE_LAB_URL, data={ "lab_code": lab.lab_code, "tidb_environment": TiDBEnvironmentType.TIDB_CLOUD_STARTER.value, "region_name": region_name, "cloud_provider": cloud_provider, }, ) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertIn("region_name", response.data) def test_create_lab_with_invalid_cloud_provider(self): lab = self._get_tidbcloud_starter_lab() region_name = lab.regions.get_all_regions()[0].value cloud_provider = "invalid_cloud_provider" response = self.client.post( self.CREATE_LAB_URL, data={ "lab_code": lab.lab_code, "tidb_environment": TiDBEnvironmentType.TIDB_CLOUD_STARTER.value, "region_name": region_name, "cloud_provider": cloud_provider, }, ) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertIn("cloud_provider", response.data) def test_create_tidbcloud_starter_lab_with_invalid_tidb_version(self): lab = self._get_tidbcloud_starter_lab() cloud_provider = lab.regions.get_supported_cloud_providers()[0] region_name = lab.regions.get_supported_regions(cloud_provider)[0].value response = self.client.post( self.CREATE_LAB_URL, data={ "lab_code": lab.lab_code, "tidb_environment": TiDBEnvironmentType.TIDB_CLOUD_STARTER.value, "region_name": region_name, "cloud_provider": cloud_provider, "tidb_version": "v7.5.1", }, ) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertIn("tidb_version", response.data) def test_create_self_managed_tidb_lab_with_invalid_tidb_version(self): lab = self._get_self_managed_tidb_lab() cloud_provider = lab.regions.get_supported_cloud_providers()[0] region_name = lab.regions.get_supported_regions(cloud_provider)[0].value tidb_version = "v6.1.1" self.assertNotIn(tidb_version, lab.tidb_versions) response = self.client.post( self.CREATE_LAB_URL, data={ "lab_code": lab.lab_code, "tidb_environment": TiDBEnvironmentType.TIDB_SELF_MANAGED.value, "region_name": region_name, "cloud_provider": cloud_provider, "tidb_version": tidb_version, }, ) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertIn("tidb_version", response.data) def test_create_self_managed_tidb_lab_with_missing_tidb_version(self): lab = self._get_self_managed_tidb_lab() cloud_provider = lab.regions.get_supported_cloud_providers()[0] region_name = lab.regions.get_supported_regions(cloud_provider)[0].value response = self.client.post( self.CREATE_LAB_URL, data={ "lab_code": lab.lab_code, "tidb_environment": TiDBEnvironmentType.TIDB_SELF_MANAGED.value, "region_name": region_name, "cloud_provider": cloud_provider, }, ) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertIn("tidb_version", response.data) def test_create_tidbcloud_starter_lab_with_tidb_version(self): lab = self._get_tidbcloud_starter_lab() cloud_provider = lab.regions.get_supported_cloud_providers()[0] region_name = lab.regions.get_supported_regions(cloud_provider)[0].value tidb_version = "v8.1.1" response = self.client.post( self.CREATE_LAB_URL, data={ "lab_code": lab.lab_code, "tidb_environment": TiDBEnvironmentType.TIDB_CLOUD_STARTER.value, "region_name": region_name, "cloud_provider": cloud_provider, "tidb_version": tidb_version, }, ) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertIn("tidb_version", response.data) self.assertIn( "This lab does not support specifying the tidb_version", str(response.data["tidb_version"][0]), ) def test_create_tidbcloud_starter_lab_without_redeeming_lab(self): lab = self._get_tidbcloud_starter_lab() cloud_provider = lab.regions.get_supported_cloud_providers()[0] region_name = lab.regions.get_supported_regions(cloud_provider)[0].value response = self.client.post( self.CREATE_LAB_URL, data={ "lab_code": lab.lab_code, "tidb_environment": TiDBEnvironmentType.TIDB_CLOUD_STARTER.value, "region_name": region_name, "cloud_provider": cloud_provider, }, ) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertIn("non_field_errors", response.data) self.assertIn("Redeem", str(response.data["non_field_errors"][0])) def test_create_tidbcloud_starter_lab_success(self): lab = self._get_tidbcloud_starter_lab() response = self.client.post(self.REDEEM_LAB_URL.format(lab_code=lab.lab_code)) self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT) cloud_provider = lab.regions.get_supported_cloud_providers()[0] region_name = lab.regions.get_supported_regions(cloud_provider)[0].value response = self.client.post( self.CREATE_LAB_URL, data={ "lab_code": lab.lab_code, "tidb_environment": TiDBEnvironmentType.TIDB_CLOUD_STARTER.value, "region_name": region_name, "cloud_provider": cloud_provider, }, ) self.assertIn("lab_instance_id", response.data) self.assertEqual(response.status_code, status.HTTP_200_OK) lab_instance = LabInstance.objects.get(lab_code=lab.lab_code, user=self.user) self.assertEqual(lab_instance.tidb_version, None) self.assertEqual( lab_instance.tidb_environment, TiDBEnvironmentType.TIDB_CLOUD_STARTER.value ) self.assertEqual(lab_instance.region_name, region_name) self.assertEqual(lab_instance.cloud_provider, cloud_provider) self.assertEqual(lab_instance.user_serverless_cluster, None) self.assertEqual(lab_instance.access_type, LabInstance.LabAccessType.CREDITS) def test_create_self_managed_tidb_lab_success(self): lab = self._get_self_managed_tidb_lab() response = self.client.post(self.REDEEM_LAB_URL.format(lab_code=lab.lab_code)) self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT) cloud_provider = lab.regions.get_supported_cloud_providers()[0] region_name = lab.regions.get_supported_regions(cloud_provider)[0].value tidb_version = lab.tidb_versions[0] response = self.client.post( self.CREATE_LAB_URL, data={ "lab_code": lab.lab_code, "tidb_environment": TiDBEnvironmentType.TIDB_SELF_MANAGED.value, "region_name": region_name, "cloud_provider": cloud_provider, "tidb_version": tidb_version, }, ) self.assertEqual(response.status_code, status.HTTP_200_OK) lab_instance = LabInstance.objects.get(lab_code=lab.lab_code, user=self.user) self.assertEqual(lab_instance.tidb_version, tidb_version) self.assertEqual( lab_instance.tidb_environment, TiDBEnvironmentType.TIDB_SELF_MANAGED.value ) self.assertEqual(lab_instance.region_name, region_name) self.assertEqual(lab_instance.cloud_provider, cloud_provider) self.assertEqual(lab_instance.user_serverless_cluster, None) self.assertEqual(lab_instance.access_type, LabInstance.LabAccessType.CREDITS)