switching to already defined setup function for testing

Author: brandonistfan

tcf_website/management/commands/fetch_enrollment.py

@@ -76,14 +76,14 @@ def fetch_section_data(section):
         # Fetch and validate response
         response = session.get(url, timeout=TIMEOUT)
         response.raise_for_status()
-        
+
         data = response.json()
 
         # Update enrollment data if available
         if data and "classes" in data and data["classes"]:
             class_data = data["classes"][0]
             section_enrollment, _ = SectionEnrollment.objects.get_or_create(section=section)
-            
+
             # Update enrollment and waitlist numbers
             section_enrollment.enrollment_taken = class_data.get("enrollment_total", 0)
             section_enrollment.enrollment_limit = class_data.get("class_capacity", 0)
@@ -93,7 +93,7 @@ def fetch_section_data(section):
 
             print(format_enrollment_update_message(section, section_enrollment))
             return True
-            
+
     except requests.exceptions.HTTPError as http_err:
         if http_err.response.status_code == 429:
             print(f"Rate limited while fetching section {section.sis_section_number}. Retrying...")
@@ -134,11 +134,11 @@ def handle(self, *args, **options):
             semester = Semester.latest()
 
         print(f"Fetching enrollment data for semester {semester}...")
-        
+
         sections = Section.objects.filter(semester=semester)
         total_sections = sections.count()
         print(f"Found {total_sections} sections")
-        
+
         # Process sections in parallel using ThreadPoolExecutor
         success_count = 0
         with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
@@ -153,7 +153,7 @@ def handle(self, *args, **options):
                 print("\nProcess interrupted by user. Shutting down...")
                 executor.shutdown(wait=False)
                 return
-        
+
         elapsed_time = time.time() - start_time
         print(f"\nSuccessfully updated {success_count}/{total_sections} sections")
         print(f"Total time: {elapsed_time:.2f} seconds")

tcf_website/tests/test_enrollment.py

@@ -1,42 +1,19 @@
 """Tests for enrollment data updates."""
 from unittest.mock import patch
 from django.test import TestCase
-from tcf_website.models import (
-    Course, Department, School, Section,
-    SectionEnrollment, Semester, Subdepartment
-)
+from tcf_website.models import SectionEnrollment
 from tcf_website.api.enrollment import update_enrollment_data
+from .test_utils import setup
 
 
 class TestEnrollment(TestCase):
     """Test cases for enrollment data updates."""
 
     def setUp(self):
         """Set up test data."""
-        self.school = School.objects.create(name="Test School")
-        self.department = Department.objects.create(
-            name="Test Department",
-            school=self.school
-        )
-        self.subdepartment = Subdepartment.objects.create(
-            name="Test Subdepartment",
-            department=self.department
-        )
-        self.semester = Semester.objects.create(
-            number="1234",
-            year=2025
-        )
-        self.course = Course.objects.create(
-            title="Test Course",
-            number=101,
-            subdepartment=self.subdepartment,
-            semester_last_taught=self.semester
-        )
-        self.section = Section.objects.create(
-            course=self.course,
-            semester=self.semester,
-            sis_section_number="99999"
-        )
+        setup(self)
+        # pylint: disable=no-member
+        self.section = self.section_course
         self.enrollment = SectionEnrollment.objects.create(
             section=self.section,
             enrollment_taken=10,
@@ -54,7 +31,7 @@ def test_mocked_update_enrollment_data(self, mock_fetch):
             "waitlist_taken": 5,
             "waitlist_limit": 10,
         }
-        update_enrollment_data(self.course.id)
+        update_enrollment_data(self.section.course.id)
         self.enrollment.refresh_from_db()
         self.assertEqual(self.enrollment.enrollment_taken, 10)
         self.assertEqual(self.enrollment.enrollment_limit, 20)
@@ -63,7 +40,7 @@ def test_mocked_update_enrollment_data(self, mock_fetch):
 
     def test_update_enrollment_data(self):
         """Test enrollment data update with real data."""
-        update_enrollment_data(self.course.id)
+        update_enrollment_data(self.section.course.id)
         self.enrollment.refresh_from_db()
         self.assertGreaterEqual(self.enrollment.enrollment_taken, 0)
         self.assertGreaterEqual(self.enrollment.enrollment_limit, 0)

tcf_website/tests/test_fetch_enrollment.py

@@ -2,42 +2,19 @@
 
 from unittest.mock import patch
 from django.test import TestCase
-from tcf_website.models import (
-    Course, Department, School, Section,
-    SectionEnrollment, Semester, Subdepartment
-)
+from tcf_website.models import SectionEnrollment
 from tcf_website.management.commands.fetch_enrollment import fetch_section_data
+from .test_utils import setup
 
 
 class FetchEnrollmentTestCase(TestCase):
     """Tests for the fetch_enrollment command."""
 
     def setUp(self):
         """Set up test data."""
-        self.school = School.objects.create(name="Test School")
-        self.department = Department.objects.create(
-            name="Test Department",
-            school=self.school
-        )
-        self.subdepartment = Subdepartment.objects.create(
-            name="Test Subdepartment",
-            department=self.department
-        )
-        self.semester = Semester.objects.create(
-            number="1234",
-            year=2025
-        )
-        self.course = Course.objects.create(
-            title="Test Course",
-            number=101,
-            subdepartment=self.subdepartment,
-            semester_last_taught=self.semester
-        )
-        self.section = Section.objects.create(
-            course=self.course,
-            semester=self.semester,
-            sis_section_number="99999"
-        )
+        setup(self)
+        # pylint: disable=no-member
+        self.section = self.section_course
 
     @patch('tcf_website.management.commands.fetch_enrollment.session.get')
     def test_fetch_enrollment_success(self, mock_get):