diff --git a/src/curriculum.py b/src/curriculum.py index 5dc8841889..2d9ceea68f 100644 --- a/src/curriculum.py +++ b/src/curriculum.py @@ -39,9 +39,13 @@ JW_SCHEDULE_TABLE_URL = "https://jw.ustc.edu.cn/ws/schedule-table/datum" MIN_CATALOG_LESSON_SEMESTER_ID = 221 MIN_CATALOG_EXAM_SEMESTER_ID = 381 +JW_SCHEDULE_CHUNK_SIZE = 100 +JW_SCHEDULE_EXPECTED_CHUNK_COUNT_KEY_PREFIX = "jw_schedule_expected_chunk_count_" -def _course_chunks(courses: list[Course], chunk_size: int = 100) -> list[list[Course]]: +def _course_chunks( + courses: list[Course], chunk_size: int = JW_SCHEDULE_CHUNK_SIZE +) -> list[list[Course]]: return [courses[i : i + chunk_size] for i in range(0, len(courses), chunk_size)] @@ -128,24 +132,109 @@ def _has_cached_catalog_exams(store: SQLiteModelStore, semester_id: str) -> bool ) -def _has_cached_jw_schedule(store: SQLiteModelStore, semester_id: str) -> bool: +def _jw_schedule_expected_chunk_count_key(semester_id: str) -> str: + return f"{JW_SCHEDULE_EXPECTED_CHUNK_COUNT_KEY_PREFIX}{semester_id}" + + +def _catalog_lesson_chunk_count( + store: SQLiteModelStore, semester_id: str +) -> int | None: + table_exists = store.conn.execute( + """ + SELECT 1 FROM sqlite_master + WHERE type = 'table' + AND name = 'catalog_teach_lesson_list_for_teach' + """ + ).fetchone() + if table_exists is None: + return None + + fetch = store.conn.execute( + """ + SELECT id FROM upstream_fetches + WHERE source = 'catalog_teach_lesson_list_for_teach' + AND ok = 1 + AND context = ? + ORDER BY id DESC + LIMIT 1 + """, + (f"semester_id={semester_id}",), + ).fetchone() + if fetch is None: + return None + + lesson_count = store.conn.execute( + """ + SELECT COUNT(*) FROM catalog_teach_lesson_list_for_teach + WHERE fetch_id = ? + """, + (fetch[0],), + ).fetchone()[0] return ( - store.conn.execute( - """ - SELECT 1 FROM upstream_fetches - WHERE source = ? - AND ok = 1 - AND (context = ? OR context LIKE ?) - LIMIT 1 - """, - ( - "jw_ws_schedule_table_datum", - f"semester_id={semester_id}", - f"%&semester_id={semester_id}", - ), - ).fetchone() - is not None + int(lesson_count) + JW_SCHEDULE_CHUNK_SIZE - 1 + ) // JW_SCHEDULE_CHUNK_SIZE + + +def _expected_jw_schedule_chunk_count( + store: SQLiteModelStore, semester_id: str +) -> int | None: + row = store.conn.execute( + "SELECT value FROM metadata WHERE key = ?", + (_jw_schedule_expected_chunk_count_key(semester_id),), + ).fetchone() + recorded_count = None + if row is not None: + try: + recorded_count = int(row[0]) + except ValueError: + return None + if recorded_count < 0: + return None + + catalog_count = _catalog_lesson_chunk_count(store, semester_id) + if catalog_count is None: + return recorded_count + if recorded_count is not None and recorded_count != catalog_count: + return None + return catalog_count + + +def _fetch_context_values(context: str | None) -> dict[str, str]: + return { + key: value + for item in (context or "").split("&") + if "=" in item + for key, value in [item.split("=", 1)] + } + + +def _has_cached_jw_schedule(store: SQLiteModelStore, semester_id: str) -> bool: + expected_count = _expected_jw_schedule_chunk_count(store, semester_id) + if expected_count is None: + return False + + successful_chunks: set[int] = set() + fetches = store.conn.execute( + """ + SELECT ok, context FROM upstream_fetches + WHERE source = 'jw_ws_schedule_table_datum' + """ ) + for ok, context in fetches: + values = _fetch_context_values(context) + if values.get("semester_id") != semester_id: + continue + if not ok: + return False + try: + chunk_index = int(values["chunk_index"]) + except (KeyError, ValueError): + return False + if chunk_index in successful_chunks: + return False + successful_chunks.add(chunk_index) + + return successful_chunks == set(range(expected_count)) def _has_cached_source_semester( @@ -330,13 +419,12 @@ async def _store_jw_schedule_chunks( courses: list[Course], ) -> None: chunks = _course_chunks(courses) - if not chunks: - guesses.add_teacher_section_guesses( - semester_id=semester_id, - catalog_lessons=catalog_response, - jw_schedules=None, - ) - return + store.put_metadata( + { + "jw_schedule_chunk_size": JW_SCHEDULE_CHUNK_SIZE, + _jw_schedule_expected_chunk_count_key(semester_id): len(chunks), + } + ) schedule_responses: list[JwWsScheduleTableDatumResponse] = [] for chunk_index, chunk in enumerate(chunks): @@ -354,13 +442,13 @@ async def _store_jw_schedule_chunks( ok=False, error=str(e), ) - logger.info( - "Skipping remaining JW schedule table chunks for semester %s after " - "non-JSON response at chunk %s", + logger.error( + "Aborting JW schedule table fetch for semester %s after non-JSON " + "response at chunk %s", semester_id, chunk_index, ) - break + raise response = JwWsScheduleTableDatumResponse.model_validate(payload) schedule_responses.append(response) @@ -377,6 +465,10 @@ async def _store_jw_schedule_chunks( context={"semester_id": semester_id, "chunk_index": chunk_index}, ) + if not _has_cached_jw_schedule(store, semester_id): + raise RuntimeError( + f"Incomplete JW schedule table chunks for semester {semester_id}" + ) guesses.add_teacher_section_guesses( semester_id=semester_id, catalog_lessons=catalog_response, diff --git a/tests/test_curriculum.py b/tests/test_curriculum.py index 7c3d76237c..1098bf211f 100644 --- a/tests/test_curriculum.py +++ b/tests/test_curriculum.py @@ -1,13 +1,21 @@ import unittest +from json import JSONDecodeError +from unittest.mock import AsyncMock, MagicMock, patch from src.curriculum import ( _cached_complete_semester_ids, + _has_cached_jw_schedule, + _jw_schedule_expected_chunk_count_key, _refresh_curriculum_semesters, _selected_curriculum_semesters, _semester_has_ended, _should_fetch_catalog_exams, _should_fetch_catalog_lessons, _should_fetch_jw_schedule_table, + _store_jw_schedule_chunks, +) +from src.models.api.catalog_api_teach_lesson_list_for_teach import ( + TeachLessonListResponse, ) from src.models.semester import Semester from src.sqlite_store import SQLiteModelStore @@ -57,6 +65,109 @@ def test_fetches_schedule_for_non_numeric_semester_ids(self) -> None: self.assertTrue(_should_fetch_jw_schedule_table("latest")) +class JwScheduleChunkTest(unittest.IsolatedAsyncioTestCase): + async def test_records_expected_count_and_accepts_all_successful_chunks( + self, + ) -> None: + store = SQLiteModelStore(":memory:") + guesses = MagicMock() + try: + catalog_fetch_id = store.record_fetch( + source="catalog_teach_lesson_list_for_teach", + method="GET", + url="lesson/401", + context={"semester_id": "401"}, + ) + store.conn.execute( + "CREATE TABLE catalog_teach_lesson_list_for_teach(" + "store_id INTEGER PRIMARY KEY AUTOINCREMENT, " + "fetch_id INTEGER NOT NULL)" + ) + store.conn.executemany( + "INSERT INTO catalog_teach_lesson_list_for_teach(fetch_id) VALUES(?)", + [(catalog_fetch_id,)] * 101, + ) + with patch( + "src.curriculum.fetch_jw_schedule_table_json", + new_callable=AsyncMock, + return_value={"result": None}, + ): + await _store_jw_schedule_chunks( + session=MagicMock(), + store=store, + guesses=guesses, + semester_id="401", + catalog_response=TeachLessonListResponse(root=[]), + courses=[MagicMock() for _ in range(101)], + ) + + metadata_key = _jw_schedule_expected_chunk_count_key("401") + expected_count = store.conn.execute( + "SELECT value FROM metadata WHERE key = ?", + (metadata_key,), + ).fetchone() + chunk_size = store.conn.execute( + "SELECT value FROM metadata WHERE key = 'jw_schedule_chunk_size'" + ).fetchone() + complete = _has_cached_jw_schedule(store, "401") + store.conn.execute("DELETE FROM metadata WHERE key = ?", (metadata_key,)) + legacy_complete = _has_cached_jw_schedule(store, "401") + finally: + store.close() + + self.assertEqual(expected_count, ("2",)) + self.assertEqual(chunk_size, ("100",)) + self.assertTrue(complete) + self.assertTrue(legacy_complete) + + async def test_missing_chunk_is_not_complete(self) -> None: + store = SQLiteModelStore(":memory:") + try: + store.put_metadata({_jw_schedule_expected_chunk_count_key("401"): 2}) + store.record_fetch( + source="jw_ws_schedule_table_datum", + method="POST", + url="jw", + context={"semester_id": "401", "chunk_index": 0}, + ) + + complete = _has_cached_jw_schedule(store, "401") + finally: + store.close() + + self.assertFalse(complete) + + async def test_failed_chunk_aborts_refresh_and_is_not_complete(self) -> None: + store = SQLiteModelStore(":memory:") + guesses = MagicMock() + try: + with ( + patch( + "src.curriculum.fetch_jw_schedule_table_json", + new_callable=AsyncMock, + side_effect=[ + {"result": None}, + JSONDecodeError("non-json", "", 0), + ], + ), + self.assertRaises(JSONDecodeError), + ): + await _store_jw_schedule_chunks( + session=MagicMock(), + store=store, + guesses=guesses, + semester_id="401", + catalog_response=TeachLessonListResponse(root=[]), + courses=[MagicMock() for _ in range(101)], + ) + + complete = _has_cached_jw_schedule(store, "401") + finally: + store.close() + + self.assertFalse(complete) + + class CatalogExamFetchTest(unittest.TestCase): def test_skips_semesters_below_minimum_exam_id(self) -> None: self.assertFalse(_should_fetch_catalog_exams("221")) @@ -153,6 +264,10 @@ def test_cached_complete_semester_ids_require_lesson_jw_and_exam_when_needed( ok=False, error="non-json", ) + for semester_id in ("221", "381", "401", "421"): + store.put_metadata( + {_jw_schedule_expected_chunk_count_key(semester_id): 1} + ) cached = _cached_complete_semester_ids( store,