fix: reject incomplete JW schedule chunks

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,

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", "<html>", 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,