From b2a6f8115079d6915bfdb264de58f8d89937c768 Mon Sep 17 00:00:00 2001
From: 0x29a <demid@opencraft.com>
Date: Tue, 7 Feb 2023 18:28:18 +0100
Subject: [PATCH] test: EnterpriseCourseEnrollmentFilterBackend

---
 test_utils/factories.py                   | 37 ++++++++++++
 tests/test_enterprise/api/test_filters.py | 70 ++++++++++++++++++++++-
 2 files changed, 106 insertions(+), 1 deletion(-)

diff --git a/test_utils/factories.py b/test_utils/factories.py
index 7b5beba7f3..bdb2f6a2fb 100644
--- a/test_utils/factories.py
+++ b/test_utils/factories.py
@@ -24,6 +24,8 @@
     EnterpriseCustomerInviteKey,
     EnterpriseCustomerReportingConfiguration,
     EnterpriseCustomerUser,
+    EnterpriseFeatureRole,
+    EnterpriseFeatureUserRoleAssignment,
     LearnerCreditEnterpriseCourseEnrollment,
     LicensedEnterpriseCourseEnrollment,
     PendingEnrollment,
@@ -264,6 +266,41 @@ class Meta:
     date_joined = factory.LazyAttribute(lambda x: FAKER.date_time_this_year(tzinfo=timezone.utc))
 
 
+class EnterpriseFeatureRoleFactory(factory.django.DjangoModelFactory):
+    """
+    EnterpriseFeatureRole factory.
+
+    Creates an instance of EnterpriseFeatureRole with minimal boilerplate.
+    """
+
+    class Meta:
+        """
+        Meta for EnterpriseFeatureRoleFactory.
+        """
+
+        model = EnterpriseFeatureRole
+
+    name = factory.LazyAttribute(lambda x: FAKER.word())
+
+
+class EnterpriseFeatureUserRoleAssignmentFactory(factory.django.DjangoModelFactory):
+    """
+    EnterpriseFeatureUserRoleAssignment factory.
+
+    Creates an instance of EnterpriseFeatureUserRoleAssignment with minimal boilerplate.
+    """
+
+    class Meta:
+        """
+        Meta for EnterpriseFeatureUserRoleAssignmentFactory.
+        """
+
+        model = EnterpriseFeatureUserRoleAssignment
+
+    role = factory.SubFactory(EnterpriseFeatureRoleFactory)
+    user = factory.SubFactory(UserFactory)
+
+
 class AnonymousUserFactory(factory.Factory):
     """
     Anonymous User factory.
diff --git a/tests/test_enterprise/api/test_filters.py b/tests/test_enterprise/api/test_filters.py
index c4a35d1510..ecb67a9ff7 100644
--- a/tests/test_enterprise/api/test_filters.py
+++ b/tests/test_enterprise/api/test_filters.py
@@ -10,7 +10,8 @@
 
 from django.conf import settings
 
-from enterprise.constants import ENTERPRISE_ADMIN_ROLE
+from enterprise.constants import ENTERPRISE_ADMIN_ROLE, ENTERPRISE_ENROLLMENT_API_ADMIN_ROLE
+from enterprise.models import EnterpriseFeatureRole
 from test_utils import FAKE_UUIDS, TEST_EMAIL, TEST_USERNAME, APITest, factories
 
 ENTERPRISE_CUSTOMER_LIST_ENDPOINT = reverse('enterprise-customer-list')
@@ -80,6 +81,73 @@ def test_filter_for_detail(self, is_staff, is_linked, expected_content_in_respon
             assert data[key] == value
 
 
+@ddt.ddt
+@mark.django_db
+class TestEnterpriseCourseEnrollmentFilterBackend(APITest):
+    """
+    Test suite for the ``EnterpriseCourseEnrollmentFilterBackend`` filter.
+    """
+
+    def setUp(self):
+        super().setUp()
+
+        self._setup_enterprise_customer_and_enrollments(
+            uuid=FAKE_UUIDS[0],
+            users=[self.user, factories.UserFactory()]
+        )
+        self._setup_enterprise_customer_and_enrollments(
+            uuid=FAKE_UUIDS[1],
+            users=[factories.UserFactory(), factories.UserFactory()]
+        )
+
+        self.url = settings.TEST_SERVER + reverse('enterprise-course-enrollment-list')
+
+    def _setup_enterprise_customer_and_enrollments(self, uuid, users):
+        """
+        Creates an enterprise customer with the uuid and enrolls passed users.
+        """
+        enterprise_customer = factories.EnterpriseCustomerFactory(uuid=uuid)
+
+        for user in users:
+            enterprise_customer_user = factories.EnterpriseCustomerUserFactory(
+                enterprise_customer=enterprise_customer,
+                user_id=user.id
+            )
+            factories.EnterpriseCourseEnrollmentFactory(
+                enterprise_customer_user=enterprise_customer_user
+            )
+
+    def _setup_user_privileges_by_role(self, user, role):
+        """
+        Sets up privileges for the passed user based on the role.
+        """
+        if role == "staff":
+            user.is_staff = True
+            user.save()
+        elif role == "enrollment_api_admin":
+            factories.EnterpriseFeatureUserRoleAssignmentFactory(
+                user=user,
+                role=EnterpriseFeatureRole.objects.get(name=ENTERPRISE_ENROLLMENT_API_ADMIN_ROLE)
+            )
+
+    @ddt.data(
+        ("regular", 1),
+        ("enrollment_api_admin", 2),
+        ("staff", 4),
+    )
+    @ddt.unpack
+    def test_filter_for_list(self, user_role, expected_course_enrollment_count):
+        """
+        Filter objects based off whether the user is a staff, enterprise enrollment api admin, or neither.
+        """
+        self._setup_user_privileges_by_role(self.user, user_role)
+
+        response = self.client.get(self.url)
+        assert response.status_code == status.HTTP_200_OK
+        data = self.load_json(response.content)
+        assert len(data['results']) == expected_course_enrollment_count
+
+
 @ddt.ddt
 @mark.django_db
 class TestEnterpriseCustomerUserFilterBackend(APITest):