fix user queryset

codeforlife/tests/api.py

@@ -76,10 +76,11 @@ def login(self, **credentials):
 
         return user
 
-    def login_teacher(self, **credentials):
+    def login_teacher(self, is_admin: bool, **credentials):
         user = self.login(**credentials)
         assert user.teacher
         assert user.teacher.school
+        assert is_admin == user.teacher.is_admin
         return user
 
     def login_student(self, **credentials):
@@ -204,6 +205,7 @@ def get_another_school_user(
         other_users: QuerySet[User],
         is_teacher: bool,
         same_school: bool,
+        same_class: bool = None,
     ):
         """
         Get a different user that is also in a school.
@@ -227,5 +229,41 @@ def get_another_school_user(
         )
         assert other_school
 
-        assert school == other_school if same_school else school != other_school
+        if same_school:
+            assert school == other_school
+
+            # Cannot assert that 2 teachers are in the same class since a class
+            # can only have 1 teacher.
+            if not (user.teacher is not None and is_teacher):
+                # At this point, same_class needs to be set.
+                assert same_class is not None, "same_class must be set."
+
+                # If one of the users is a teacher.
+                if user.teacher is not None or is_teacher:
+                    # Get the teacher.
+                    teacher = other_user if is_teacher else user
+
+                    # Get the student's class' teacher.
+                    class_teacher = (
+                        user if is_teacher else other_user
+                    ).student.class_field.teacher.new_user
+
+                    # Assert the teacher is the class' teacher.
+                    assert (
+                        teacher == class_teacher
+                        if same_class
+                        else teacher != class_teacher
+                    )
+                # Else, both users are students.
+                else:
+                    assert (
+                        user.student.class_field
+                        == other_user.student.class_field
+                        if same_class
+                        else user.student.class_field
+                        != other_user.student.class_field
+                    )
+        else:
+            assert school != other_school
+
         return other_user

codeforlife/user/tests/views/test_school.py

@@ -65,6 +65,7 @@ def _login_teacher(self):
         return self.client.login_teacher(
             email="[email protected]",
             password="Password1",
+            is_admin=True,
         )
 
     def _login_student(self):

codeforlife/user/tests/views/test_user.py

@@ -62,9 +62,17 @@ def setUp(self):
         )
 
     def _login_teacher(self):
+        return self.client.login_teacher(
+            email="[email protected]",
+            password="Password1",
+            is_admin=False,
+        )
+
+    def _login_admin_teacher(self):
         return self.client.login_teacher(
             email="[email protected]",
             password="Password1",
+            is_admin=True,
         )
 
     def _login_student(self):
@@ -81,10 +89,11 @@ def _login_indy_student(self):
 
     """
     Retrieve naming convention:
-        test_retrieve__{user_type}__{other_user_type}__{same_school}
+        test_retrieve__{user_type}__{other_user_type}__{same_school}__{same_class}
 
     user_type: The type of user that is making the request. Options:
-        - teacher: A teacher.
+        - teacher: A non-admin teacher.
+        - admin_teacher: An admin teacher.
         - student: A school student.
         - indy_student: A non-school student.
 
@@ -96,6 +105,10 @@ def _login_indy_student(self):
     same_school: A flag for if the other user is from the same school. Options:
         - same_school: The other user is from the same school.
         - not_same_school: The other user is not from the same school.
+
+    same_class: A flag for if the other user is from the same class. Options:
+        - same_class: The other user is from the same class.
+        - not_same_class: The other user is not from the same class.
     """
 
     def _retrieve_user(
@@ -155,9 +168,30 @@ def test_retrieve__teacher__teacher__same_school(self):
 
         self._retrieve_user(other_user)
 
-    def test_retrieve__teacher__student__same_school(self):
+    def test_retrieve__teacher__student__same_school__same_class(self):
         """
-        Teacher can retrieve a student from the same school.
+        Teacher can retrieve a student from the same school and class.
+        """
+
+        user = self._login_teacher()
+
+        other_user = self.get_another_school_user(
+            user,
+            other_users=User.objects.filter(
+                new_student__class_field__teacher__school=user.teacher.school,
+                new_student__class_field__teacher=user.teacher,
+            ),
+            is_teacher=False,
+            same_school=True,
+            same_class=True,
+        )
+
+        self._retrieve_user(other_user)
+
+    def test_retrieve__teacher__student__same_school__not_same_class(self):
+        """
+        Teacher cannot retrieve a student from the same school and a different
+        class.
         """
 
         user = self._login_teacher()
@@ -166,48 +200,143 @@ def test_retrieve__teacher__student__same_school(self):
             user,
             other_users=User.objects.filter(
                 new_student__class_field__teacher__school=user.teacher.school
+            ).exclude(new_student__class_field__teacher=user.teacher),
+            is_teacher=False,
+            same_school=True,
+            same_class=False,
+        )
+
+        self._retrieve_user(
+            other_user,
+            status_code_assertion=status.HTTP_404_NOT_FOUND,
+        )
+
+    def test_retrieve__admin_teacher__student__same_school__same_class(self):
+        """
+        Admin teacher can retrieve a student from the same school and class.
+        """
+
+        user = self._login_admin_teacher()
+
+        other_user = self.get_another_school_user(
+            user,
+            other_users=User.objects.filter(
+                new_student__class_field__teacher__school=user.teacher.school,
+                new_student__class_field__teacher=user.teacher,
             ),
             is_teacher=False,
             same_school=True,
+            same_class=True,
+        )
+
+        self._retrieve_user(other_user)
+
+    def test_retrieve__admin_teacher__student__same_school__not_same_class(
+        self,
+    ):
+        """
+        Admin teacher can retrieve a student from the same school and a
+        different class.
+        """
+
+        user = self._login_admin_teacher()
+
+        other_user = self.get_another_school_user(
+            user,
+            other_users=User.objects.filter(
+                new_student__class_field__teacher__school=user.teacher.school
+            ).exclude(new_student__class_field__teacher=user.teacher),
+            is_teacher=False,
+            same_school=True,
+            same_class=False,
         )
 
         self._retrieve_user(other_user)
 
-    def test_retrieve__student__teacher__same_school(self):
+    def test_retrieve__student__teacher__same_school__same_class(self):
         """
-        Student cannot retrieve a teacher from the same school.
+        Student cannot retrieve a teacher from the same school and class.
         """
 
         user = self._login_student()
 
         other_user = self.get_another_school_user(
             user,
             other_users=User.objects.filter(
-                new_teacher__school=user.student.class_field.teacher.school
+                new_teacher__school=user.student.class_field.teacher.school,
+                new_teacher__class_teacher=user.student.class_field,
             ),
             is_teacher=True,
             same_school=True,
+            same_class=True,
         )
 
         self._retrieve_user(
             other_user,
             status_code_assertion=status.HTTP_404_NOT_FOUND,
         )
 
-    def test_retrieve__student__student__same_school(self):
+    def test_retrieve__student__teacher__same_school__not_same_class(self):
         """
-        Student cannot retrieve another student from the same school.
+        Student cannot retrieve a teacher from the same school and a different
+        class.
+        """
+
+        user = self._login_student()
+
+        other_user = self.get_another_school_user(
+            user,
+            other_users=User.objects.filter(
+                new_teacher__school=user.student.class_field.teacher.school
+            ).exclude(new_teacher__class_teacher=user.student.class_field),
+            is_teacher=True,
+            same_school=True,
+            same_class=False,
+        )
+
+        self._retrieve_user(
+            other_user,
+            status_code_assertion=status.HTTP_404_NOT_FOUND,
+        )
+
+    def test_retrieve__student__student__same_school__same_class(self):
+        """
+        Student can retrieve another student from the same school and class.
         """
 
         user = self._login_student()
 
         other_user = self.get_another_school_user(
             user,
             other_users=User.objects.exclude(id=user.id).filter(
-                new_student__class_field__teacher__school=user.student.class_field.teacher.school
+                new_student__class_field__teacher__school=user.student.class_field.teacher.school,
+                new_student__class_field=user.student.class_field,
             ),
             is_teacher=False,
             same_school=True,
+            same_class=True,
+        )
+
+        self._retrieve_user(other_user)
+
+    def test_retrieve__student__student__same_school__not_same_class(self):
+        """
+        Student cannot retrieve another student from the same school and a
+        different class.
+        """
+
+        user = self._login_student()
+
+        other_user = self.get_another_school_user(
+            user,
+            other_users=User.objects.exclude(id=user.id)
+            .filter(
+                new_student__class_field__teacher__school=user.student.class_field.teacher.school,
+            )
+            .exclude(new_student__class_field=user.student.class_field),
+            is_teacher=False,
+            same_school=True,
+            same_class=False,
         )
 
         self._retrieve_user(
@@ -345,7 +474,7 @@ def test_retrieve__indy_student__student(self):
         - teacher: A teacher.
         - student: A school student.
         - indy_student: A non-school student.
-    
+
     filters: Any search params used to dynamically filter the list.
     """
 
@@ -373,13 +502,14 @@ def test_list__teacher(self):
         self._list_users(
             User.objects.filter(new_teacher__school=user.teacher.school)
             | User.objects.filter(
-                new_student__class_field__teacher__school=user.teacher.school
+                new_student__class_field__teacher__school=user.teacher.school,
+                new_student__class_field__teacher=user.teacher,
             )
         )
 
     def test_list__teacher__students_in_class(self):
         """
-        Teacher can list all the users in the same school.
+        Teacher can list all the users in a class they own.
         """
 
         user = self._login_teacher()

codeforlife/user/views/user.py

@@ -13,14 +13,25 @@ class UserViewSet(ModelViewSet):
     def get_queryset(self):
         user: User = self.request.user
         if user.teacher is None:
-            return User.objects.filter(id=user.id)
+            if user.student.class_field is None:
+                return User.objects.filter(id=user.id)
+
+            return User.objects.filter(
+                new_student__class_field=user.student.class_field
+            )
 
         teachers = User.objects.filter(
             new_teacher__school=user.teacher.school_id
         )
-        students = User.objects.filter(
-            # TODO: add school foreign key to student model.
-            new_student__class_field__teacher__school=user.teacher.school_id,
+        students = (
+            User.objects.filter(
+                # TODO: add school foreign key to student model.
+                new_student__class_field__teacher__school=user.teacher.school_id,
+            )
+            if user.teacher.is_admin
+            else User.objects.filter(
+                new_student__class_field__teacher=user.teacher
+            )
         )
 
         return teachers | students