tests/test_sprint.py

"""Offline tests for the sprint's parser, retrieval and server changes."""
import unittest
import xml.etree.ElementTree as ET

from canlex.ingest import _schedule_chunks, _table_text
from canlex.index import _fold_accents, tokenize
from canlex.server import _normalize_section
from canlex.citator import canlii_url_from_citation

_SRC = {"short": "TestAct", "name": "Test Act", "web_base": "https://x/"}

_SCHED_XML = """<Statute>
<Schedule>
 <ScheduleFormHeading><Label>SCHEDULE I</Label></ScheduleFormHeading>
 <List><Item><Label>1</Label><Text>Opium (all salts)</Text>
   <List><Item><Label>(1)</Label><Text>Ketamine</Text></Item></List>
 </Item></List>
</Schedule>
<Schedule>
 <ScheduleFormHeading><Label>SCHEDULE II</Label></ScheduleFormHeading>
 <Repealed>[Repealed]</Repealed>
</Schedule>
<Schedule id="NifProvs">
 <ScheduleFormHeading><TitleText>AMENDMENTS NOT IN FORCE</TitleText></ScheduleFormHeading>
 <List><Item><Text>never rendered here</Text></Item></List>
</Schedule>
<Schedule>
 <ScheduleFormHeading><Label>SCHEDULE III</Label><TitleText>Violations</TitleText></ScheduleFormHeading>
 <TableGroup><table><tgroup><tbody>
   <row><entry>Provision</entry><entry>Penalty</entry></row>
   <row><entry>s. 12</entry><entry>$500</entry></row>
 </tbody></tgroup></table></TableGroup>
</Schedule>
</Statute>"""


class ScheduleParsingTests(unittest.TestCase):
    def setUp(self):
        self.root = ET.fromstring(_SCHED_XML)

    def chunks(self, src=None):
        return _schedule_chunks(self.root, "T-1", src or _SRC, "2026-01-01")

    def test_list_schedule_rendered_with_nested_items(self):
        c = self.chunks()[0]
        self.assertIn("Opium", c["text"])
        self.assertIn("Ketamine", c["text"])
        self.assertEqual(c["section"], "Schedule I")
        self.assertEqual(c["status"], "in force")

    def test_repealed_schedule_flagged(self):
        c = next(x for x in self.chunks() if x["section"] == "Schedule Ii")
        self.assertEqual(c["status"], "repealed")

    def test_nifprovs_excluded(self):
        self.assertFalse([c for c in self.chunks()
                          if "never rendered" in c["text"]])

    def test_table_rows_pipe_separated(self):
        c = next(x for x in self.chunks() if "Violations" in x["marginal_note"])
        self.assertIn("s. 12 | $500", c["text"])

    def test_skip_schedules_optout(self):
        self.assertEqual(self.chunks({**_SRC, "skip_schedules": True}), [])


class AccentFoldingTests(unittest.TestCase):
    def test_fold(self):
        self.assertEqual(_fold_accents("détention"), "detention")

    def test_tokenize_french(self):
        # Folded French stems to the same normalized token as the English word
        # ('detention' and 'détention' both -> 'detain').
        self.assertEqual(tokenize("détention"), tokenize("detention"))

    def test_english_unchanged(self):
        self.assertEqual(tokenize("seizure of goods"),
                         tokenize("seizure of goods"))


class SectionNormalizationTests(unittest.TestCase):
    CASES = {
        "s. 34(1)(c)": "34", "section 20.1": "20.1", "s 34": "34",
        "ss. 18(2)": "18", "34": "34", "Schedule I": "Schedule I",
        "schedule 2": "schedule 2", "§34": "34",
    }

    def test_normalization_table(self):
        for given, want in self.CASES.items():
            self.assertEqual(_normalize_section(given), want, given)


class CitatorCourtMapTests(unittest.TestCase):
    def test_fpslreb(self):
        self.assertIn("/ca/pslreb/doc/2023/2023fpslreb12/",
                      canlii_url_from_citation("2023 FPSLREB 12"))

    def test_onca_provincial_segment(self):
        self.assertIn("/on/onca/doc/2024/2024onca608/",
                      canlii_url_from_citation("2024 ONCA 608"))

    def test_scc_unchanged(self):
        self.assertIn("/ca/scc/doc/2019/2019scc65/",
                      canlii_url_from_citation("2019 SCC 65"))


class BilingualTests(unittest.TestCase):
    def test_query_lang(self):
        from canlex.index import query_lang
        self.assertEqual(query_lang("délai de contrôle de la détention"), "fr")
        self.assertEqual(query_lang("interdiction de territoire"), "fr")
        self.assertEqual(query_lang("detention review timelines"), "en")
        self.assertEqual(query_lang("seizure of goods s. 110"), "en")

    def test_frenchify(self):
        from canlex.ingest import _frenchify
        chunks = [{"id": "I-2.5-s36", "citation": "IRPA, s. 36",
                   "source_url": "https://x/eng/acts/i-2.5/"}]
        out = _frenchify(chunks, "I-2.5")
        self.assertEqual(out[0]["id"], "I-2.5-fr-s36")
        self.assertEqual(out[0]["lang"], "fr")
        self.assertEqual(out[0]["citation"], "IRPA, art. 36")
        self.assertIn("/fra/lois/", out[0]["source_url"])

    def test_fr_xml_url_dors(self):
        from canlex.ingest import _fr_xml_url
        self.assertIn("/fra/XML/DORS-2002-227.xml", _fr_xml_url(
            {"xml_url": "https://x/eng/XML/SOR-2002-227.xml"}))


class MultiVectorTests(unittest.TestCase):
    def test_embed_texts_windows(self):
        from canlex.embed import embed_texts, _MAX_BODY, _MAX_WINDOWS
        short = {"id": "x", "act_short": "A", "marginal_note": "Note",
                 "heading": "", "part": "", "text": "short body",
                 "doc_type": "legislation", "section": "1"}
        self.assertEqual(len(embed_texts(short)), 1)
        long = dict(short, text="w " * (_MAX_BODY * 12))
        texts = embed_texts(long)
        self.assertEqual(len(texts), _MAX_WINDOWS)
        # every tail window keeps the topical anchor
        self.assertTrue(all(t.startswith("A . ") for t in texts[1:]))


class TypoCorrectionTests(unittest.TestCase):
    def test_edit_distance(self):
        from canlex.index import LegislationIndex as L
        self.assertEqual(L._edit_distance("detension", "detention"), 1)
        self.assertEqual(L._edit_distance("undeclraed", "undeclared"), 1)
        self.assertEqual(L._edit_distance("serius", "serious"), 1)
        self.assertEqual(L._edit_distance("same", "same"), 0)
        self.assertGreater(L._edit_distance("apple", "orange"), 2)

    def test_edit_distance_transposition_costs_one(self):
        from canlex.index import LegislationIndex as L
        self.assertEqual(L._edit_distance("sieze", "seize"), 1)

    @staticmethod
    def _index(text, postings=None):
        """A bare index carrying just what the correction layer reads."""
        from canlex.index import LegislationIndex
        idx = LegislationIndex.__new__(LegislationIndex)
        idx.chunks = [{"text": text}]
        idx.postings = postings or {}
        idx._word_vocab = idx._word_tri = None
        return idx

    def test_corrects_to_the_nearest_corpus_word(self):
        idx = self._index("detention review " * 6)
        self.assertEqual(idx._correct_word("detension"), "detention")

    def test_rare_corpus_word_is_never_a_correction(self):
        # Frequency floor of 5: OCR debris must not win a correction.
        idx = self._index("detentlon " * 4 + "detention " * 9)
        self.assertEqual(idx._correct_word("detension"), "detention")

    def test_known_query_passes_through_untouched(self):
        query = "detention review timelines"
        idx = self._index("detention review " * 6,
                          postings={"detent": [0], "review": [0],
                                    "timelin": [0]})
        self.assertIs(idx._correct_query(query), query)

    def test_unknown_word_is_replaced(self):
        idx = self._index("detention review " * 6, postings={"review": [0]})
        self.assertEqual(idx._correct_query("detension review"),
                         "detention review")

    def test_nearly_unknown_word_is_appended_not_replaced(self):
        # df <= 2 -- a rare legitimate term keeps its own recall, so the
        # correction is added beside it rather than swapped in.
        idx = self._index("seizure " * 120 + "seizuer ",
                          postings={"seizuer": [0]})
        self.assertEqual(idx._correct_query("seizuer"), "seizuer seizure")


class LanguageScopedBM25Tests(unittest.TestCase):
    """The French twins must not weight English retrieval (see _build_bm25)."""

    @staticmethod
    def _index(chunks):
        from canlex.index import LegislationIndex
        idx = LegislationIndex.__new__(LegislationIndex)
        idx.chunks = chunks
        idx._tri_index = None
        idx._build_bm25()
        return idx

    @staticmethod
    def _chunk(text, lang="en"):
        return {"doc_type": "legislation", "act_code": "I-2.5", "act_name": "A",
                "act_short": "A", "section": "1", "marginal_note": "",
                "heading": "", "part": "", "division": "", "lang": lang,
                "text": text}

    def test_french_chunks_do_not_move_english_idf(self):
        en = [self._chunk("detention review"), self._chunk("seizure")]
        idx_en = self._index(list(en))
        idx_mixed = self._index(en + [self._chunk("detention", lang="fr")] * 8)
        self.assertAlmostEqual(idx_en.idf["detain"], idx_mixed.idf_en["detain"])
        # ... while the whole-corpus statistic, which French queries use, does
        # see them.
        self.assertNotAlmostEqual(idx_en.idf["detain"], idx_mixed.idf["detain"])

    def test_french_chunks_do_not_move_english_avgdl(self):
        en = [self._chunk("detention review of a permanent resident")]
        idx_en = self._index(list(en))
        idx_mixed = self._index(en + [self._chunk("a b c d e f g h", lang="fr")])
        self.assertAlmostEqual(idx_en.avgdl, idx_mixed.avgdl_en)
        self.assertNotAlmostEqual(idx_en.avgdl, idx_mixed.avgdl)

    def test_french_only_term_scores_nothing_in_english_scope(self):
        idx = self._index([self._chunk("detention review"),
                           self._chunk("annulé", lang="fr")])
        self.assertEqual(idx._bm25_scores("annulé", en_only=True), {})
        self.assertTrue(idx._bm25_scores("annulé"))


class MemorandumSourceKeyTests(unittest.TestCase):
    """The diversity cap must key each memorandum on its parent document."""

    @staticmethod
    def _keys(chunks):
        from canlex.index import LegislationIndex
        idx = LegislationIndex.__new__(LegislationIndex)
        idx.chunks = [dict(c, doc_type="memorandum") for c in chunks]
        return [idx._source_key(i) for i in range(len(idx.chunks))]

    def test_enf_sections_share_their_chapter(self):
        keys = self._keys([{"id": "enf-10-21", "act_code": "ENF-10",
                            "section": "ENF 10 s. 21"},
                           {"id": "enf-10-6", "act_code": "ENF-10",
                            "section": "ENF 10 s. 6"},
                           {"id": "enf-19-11", "act_code": "ENF-19",
                            "section": "ENF 19 s. 11"}])
        self.assertEqual(keys[0], keys[1])
        self.assertNotEqual(keys[0], keys[2])

    def test_pdi_pages_share_their_tree(self):
        keys = self._keys([{"id": "pdi-refugee-31-2", "act_code": "PDI-REFUGEE",
                            "section": "PRRA: Intake"},
                           {"id": "pdi-refugee-34-0", "act_code": "PDI-REFUGEE",
                            "section": "PRRA: Applicant"}])
        self.assertEqual(keys[0], keys[1])

    def test_flat_families_key_on_the_item(self):
        # One act_code covers every D-memo and every AMPS contravention, so
        # there the cap must bind per memo / per contravention instead.
        keys = self._keys([{"id": "dmemo-D19-9-2-1", "act_code": "D-Memo",
                            "section": "D19-9-2"},
                           {"id": "dmemo-D19-9-2-2", "act_code": "D-Memo",
                            "section": "D19-9-2"},
                           {"id": "dmemo-D2-3-1-1", "act_code": "D-Memo",
                            "section": "D2-3-1"},
                           {"id": "amps-C001", "act_code": "AMPS",
                            "section": "C001"},
                           {"id": "amps-C004", "act_code": "AMPS",
                            "section": "C004"}])
        self.assertEqual(keys[0], keys[1])
        self.assertNotEqual(keys[1], keys[2])
        self.assertNotEqual(keys[3], keys[4])


class RefreshCuratedTests(unittest.TestCase):
    def test_check_curated_runs_offline(self):
        from canlex.refresh import check_curated
        rows = check_curated("2026-07-22")
        files = {r["file"] for r in rows}
        self.assertIn("us_dispositions.json", files)
        for r in rows:
            self.assertIn(r["status"], ("ok", "aging", "drift", "error"))


if __name__ == "__main__":
    unittest.main()