Reduce flakiness of export tests

Author: MartinVanIJcken

mavis/test/data_models.py

@@ -66,15 +66,20 @@ def to_onboarding(self) -> str:
 
     @classmethod
     def get_from_testing_api(
-        cls, base_url: str, year_groups: dict[str, int]
+        cls,
+        base_url: str,
+        year_groups: dict[str, int],
+        *,
+        required_year_groups: list[int] | None = None,
     ) -> "dict[str, list[School]]":
         def _get_schools_with_year_group(year_group: int) -> list[School]:
+            filter_groups = required_year_groups if required_year_groups else [year_group]
             url = urllib.parse.urljoin(base_url, "api/testing/locations")
             params = {
                 "type": "gias_school",
                 "status": "open",
                 "is_attached_to_team": "false",
-                "gias_year_groups[]": [str(year_group)],
+                "gias_year_groups[]": [str(yg) for yg in filter_groups],
                 "site": "",
             }
 

mavis/test/onboarding.py

@@ -61,14 +61,20 @@ class PointOfCareOnboarding(Onboarding):
 
     @classmethod
     def get_onboarding_data_for_tests(
-        cls, base_url: str, year_groups: dict[str, int]
+        cls,
+        base_url: str,
+        year_groups: dict[str, int],
+        *,
+        required_year_groups: list[int] | None = None,
     ) -> "PointOfCareOnboarding":
         organisation = Organisation.generate()
         subteam = Subteam.generate()
         team = PointOfCareTeam.generate(subteam, organisation)
         users = cls._generate_users()
         clinics = [Clinic.generate()]
-        schools = School.get_from_testing_api(base_url, year_groups)
+        schools = School.get_from_testing_api(
+            base_url, year_groups, required_year_groups=required_year_groups
+        )
 
         return cls(
             organisation=organisation,

tests/test_reporting_regression.py

@@ -41,6 +41,7 @@ def _onboard_team(base_url):
     onboarding = PointOfCareOnboarding.get_onboarding_data_for_tests(
         base_url=base_url,
         year_groups=_year_groups,
+        required_year_groups=[_yg1, _yg2, _yg3],
     )
     return _create_onboarding_with_retry(base_url, onboarding)