grabber-ai/tests/run_tests.py

"""Юнит-тесты DUB-AI (запуск: .venv/Scripts/python tests/run_tests.py).

Проверяют ключевые части без обращения к GPU:
  * поиск спана автора (точно / алиас / инициалы / кавычки / разделители);
  * канонизация имён студий;
  * снятие префикса строки дорожки;
  * сегментацию страниц на дорожки;
  * HTML -> текст;
  * BIO-выравнивание спана по offsets токенизатора.
"""

from __future__ import annotations

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from src.html_extract import (RE_LINE_PREFIX, _split_entries,  # noqa: E402
                              extract_segments, html_to_text)
from src.studio_norm import canonical, find_studio_span  # noqa: E402

PASS = 0
FAIL = 0


def check(name: str, cond: bool, extra: str = "") -> None:
    global PASS, FAIL
    if cond:
        PASS += 1
        print(f"  ok  {name}")
    else:
        FAIL += 1
        print(f"FAIL  {name}  {extra}")


# ---------------------------------------------------------------- спаны
print("== find_studio_span ==")
txt = "Профессиональный (многоголосый закадровый) | HDRezka Studio 18+"
s, e = find_studio_span(txt, "HDrezka")
check("HDrezka найден", s is not None and "HDRezka" in txt[s:e], f"span={(s, e)}")

s, e = find_studio_span("Авторский (одноголосый закадровый) - Ю. Сербин", "Юрий Сербин")
check("инициалы Ю. Сербин", "Сербин" in "Авторский (одноголосый закадровый) - Ю. Сербин"[s:e] if s is not None else False)

s, e = find_studio_span("Авторский (одноголосый закадровый) Гаврилов", "Андрей Гаврилов")
check("фамилия Гаврилов", s is not None and e is not None and s > 0)

s, e = find_studio_span('Профессиональный (многоголосый закадровый) Студия «Кравец», LostFilm',
                        "Студия Кравец, LostFilm")
check("кавычки в тексте", s is not None)

s, e = find_studio_span("Профессиональный (многоголосый закадровый) - SDI Media | TV 1000",
                        "SDI Media TV 1000")
check("разделитель |", s is not None)

check("нет автора -> None", find_studio_span("Профессиональный (дублированный)", "") is None)
check("нет вхождения -> None", find_studio_span("Профессиональный (дублированный)", "HDrezka") is None)

# ---------------------------------------------------------------- канонизация
print("== canonical ==")
check("HDRezka Studio -> HDrezka", canonical("HDRezka Studio") == "HDrezka")
check("hdrezka -> HDrezka", canonical("hdrezka") == "HDrezka")
check("Read Head Sound -> Red Head Sound", canonical("Read Head Sound") == "Red Head Sound")
check("Ultradox Studio -> Ultradox", canonical("Ultradox Studio") == "Ultradox")
check("Jaskier ненормативная лексика -> Jaskier",
      canonical("Jaskier ненормативная лексика") == "Jaskier")
check("официальный -> ''", canonical("официальный") == "")

# ---------------------------------------------------------------- префиксы
print("== RE_LINE_PREFIX ==")
for line, body in [
    ("Перевод: Профессиональный (дублированный)", "Профессиональный (дублированный)"),
    ("Перевод 2: Любительский (одноголосый закадровый)", "Любительский (одноголосый закадровый)"),
    ("Перевод #1: Дублированный", "Дублированный"),
    ("Озвучивание: Дублированное | Official", "Дублированное | Official"),
    ("Аудио 1: Opus, 6 ch, 221 Kbps | Русский AVO (Дмитрий Пучков)",
     "Opus, 6 ch, 221 Kbps | Русский AVO (Дмитрий Пучков)"),
    ("Перевод:Профессиональный (дублированный)", "Профессиональный (дублированный)"),
]:
    check(f"префикс «{line[:20]}…»", RE_LINE_PREFIX.sub("", line) == body)

# ---------------------------------------------------------------- сегментация
print("== extract_segments ==")
segs = extract_segments(
    "Описание релиза\n"
    "Перевод 1: Профессиональный (дублированный)\n"
    "Перевод 2: Профессиональный (двухголосый закадровый) tycoon\n"
    "Субтитры: русские, английские\n")
check("две дорожки", len(segs) == 2, f"got {segs}")

segs = extract_segments("Перевод 1: Профессиональный (дублированный) 2. Авторский (одноголосый закадровый) Гоблин")
check("нумерация 1./2.", len(segs) == 2, f"got {segs}")

segs = extract_segments("Перевод: Одноголосый, многоголосый закадровый")
check("запятая-перечисление", len(segs) == 2, f"got {segs}")

segs = extract_segments("Перевод: Профессиональный (многоголосый, закадровый)")
check("скобки-уточнение не режутся", len(segs) == 1, f"got {segs}")

segs = extract_segments("Перевод:\nПрофессиональный (дублированный) Невафильм")
check("значение на следующей строке",
      segs == ["Профессиональный (дублированный) Невафильм"], f"got {segs}")

segs = extract_segments("Перевод: Профессиональный (двухголосый закадровый) + Любительский (одноголосый закадровый)")
check("разделитель +", len(segs) == 2, f"got {segs}")

segs = extract_segments("Аудио 1: Opus, 6 ch, 221 Kbps | Русский AVO (Дмитрий Пучков)\n"
                        "Аудио 2: Opus, 6 ch, 199 Kbps | Русский DVO (Кубик в Кубе)")
check("nnm аудио-дорожки", len(segs) == 2, f"got {segs}")

segs = extract_segments("Аудио: Русский (AC3, 6 ch, 640 Кбит/с), Английский (AC3)")
check("codec-строка без типа пропущена", len(segs) == 0, f"got {segs}")

check("_split_entries одиночная", _split_entries("Дублированный") == ["Дублированный"])

# ---------------------------------------------------------------- html
print("== html_to_text ==")
html = ("<html><body><div>меню форума</div>"
        '<div class="post_body">Перевод: Профессиональный (дублированный) Невафильм</div>'
        "</body></html>")
text = html_to_text(html)
check("селектор .post_body", "Невафильм" in text and "меню форума" not in text)

html2 = "<html><body><p>hello</p></body></html>"
check("fallback на весь текст", "hello" in html_to_text(html2))

print("== read_input / mojibake ==")
import tempfile

from src.html_extract import read_input

with tempfile.TemporaryDirectory() as td:
    p = Path(td) / "mojibake.html"
    p.write_bytes("Перевод: Дублированный".encode("utf-8").decode("cp1251",
                  errors="ignore").encode("utf-8"))
    check("mojibake восстановлен", "Дублированный" in read_input(p))
    p2 = Path(td) / "cp1251.txt"
    p2.write_bytes("Перевод: Дублированный".encode("cp1251"))
    check("cp1251 прочитан", "Дублированный" in read_input(p2))

# ---------------------------------------------------------------- BIO
print("== BIO-выравнивание (требует токенизатор) ==")
try:
    from transformers import AutoTokenizer

    from src.train import make_collate

    tok = AutoTokenizer.from_pretrained("models/best")
    collate = make_collate(tok, {"проф. дубляж": 0})
    text = "Авторский (одноголосый закадровый) - Сергей Визгунов"
    s0 = text.index("Сергей")
    enc, types, bio, ages, offsets = collate(
        [{"text": text, "translation_type": "проф. дубляж",
          "span": [s0, s0 + len("Сергей Визгунов")], "is_18plus": False}])
    marked = {i for i, lab in enumerate(bio[0].tolist()) if lab >= 1}
    span_toks = [tok.convert_ids_to_tokens(int(enc["input_ids"][0, i])) for i in sorted(marked)]
    joined = "".join(t.lstrip("#") for t in span_toks).lower()
    # у ruBert-base токенизатор strips accents: «й» может стать «и»
    check("токены спана = ФИ", "визгунов" in joined and joined.startswith("сер"),
          f"tokens={span_toks}")
    check("первый токен B", bio[0, min(marked)].item() == 1 if marked else False)
except Exception as e:  # noqa: BLE001
    print(f"  пропущено (нет модели models/best): {e}")

# ---------------------------------------------------------------- релиз / ссылки
print("== release / fetch ==")
from src.fetch import decode_bytes, detect_tracker  # noqa: E402
from src.release import aggregate, title_segments  # noqa: E402

check("tracker rutracker", detect_tracker("https://rutracker.org/forum/viewtopic.php?t=1") == "rutracker")
check("tracker kinozal", detect_tracker("https://kinozal.tv/details.php?id=1") == "kinozal")
check("tracker nnm", detect_tracker("https://nnmclub.to/forum/viewtopic.php?t=1") == "nnm")
for bad in ("http://localhost:8000/", "http://127.0.0.1/", "http://169.254.169.254/",
            "https://rutracker.org.evil.com/", "file:///etc/passwd", "ftp://rutracker.org/"):
    try:
        detect_tracker(bad)
        check(f"SSRF отклонён: {bad}", False)
    except ValueError:
        check(f"SSRF отклонён: {bad}", True)
check("decode cp1251", decode_bytes("Перевод".encode("cp1251")) == "Перевод")
check("decode charset", decode_bytes("Перевод".encode("cp1251"), "text/html; charset=windows-1251") == "Перевод")

check("title: ПМ (студия)",
      title_segments("Гангстерленд / 2025 / ПМ (LostFilm) / WEB-DLRip")
      == ["Профессиональный (многоголосый закадровый) | LostFilm"])
check("title: ДБ + список ПМ без студий",
      title_segments("Х / 2026 / ДБ (Red Head Sound), 5 x ПМ, ПД, ЛМ, СТ / 4K")
      == ["Профессиональный (дублированный) | Red Head Sound"])
check("title: пусто", title_segments("Просто фильм / 2020 / WEB-DL") == [])


def _t(typ, st, age=False):
    return {"segment": f"{typ}|{st}", "translation_type": typ, "studio": st, "is_18plus": age}


agg = aggregate([_t("проф. дубляж", "MovieDalen"), _t("проф. дубляж", "moviedalen", True),
                 _t("проф. многоголосый", "LostFilm")])
check("aggregate: дедуп", len(agg["tracks"]) == 2, str(agg))
check("aggregate: 18+ = OR", agg["is_18plus"] is True)
check("aggregate: студии", [x["studio"] for x in agg["studios"]] == ["MovieDalen", "LostFilm"])
check("aggregate: пусто", aggregate([])["is_18plus"] is False)

from src.fetch import blocked_reason  # noqa: E402
from src.release import merge_title_tracks, title_from_text  # noqa: E402

check("заглушка Cloudflare", blocked_reason("<html><title>Один момент…</title>") is not None)
check("нормальная страница", blocked_reason("<html><title>X</title>Перевод: A") is None)

_page = [dict(_t("проф. двухголосый", None), studio_span=None, studio_confidence=None),
         dict(_t("проф. дубляж", "MovieDalen"), studio_span="MovieDalen", studio_confidence=1)]
_title = [dict(_t("проф. двухголосый", "Red Head Sound"), studio_span="Red Head Sound",
               studio_confidence=0.9),
          dict(_t("проф. дубляж", "MovieDalen"), studio_span="x", studio_confidence=1),
          dict(_t("проф. многоголосый", "LostFilm"), studio_span="LostFilm", studio_confidence=1)]
merge_title_tracks(_page, _title)
check("title: студия вписана в дорожку без студии",
      _page[0]["studio"] == "Red Head Sound" and _page[0]["source"] == "page+title")
check("title: дубль студии пропущен, новая добавлена",
      len(_page) == 3 and _page[2]["source"] == "title")

from src.fetch import inline_tabs, kinozal_tabs  # noqa: E402
from src.release import codec_segments, merge_extra_tracks, same_studio  # noqa: E402

_kz = ('<ul><li><a onclick="showtab(7,100); return false;" href="#">Техданные</a></li>'
       '<li><a onclick="showtab(7,0); return false;" href="#">Релиз</a></li>'
       '<li><a onclick="showtab(7,2); return false;" href="#">Скриншоты</a></li></ul>'
       '<div id="tabs"><b>Перевод:</b> Профессиональный многоголосый</div>')
check("kinozal: вкладки без «Техданных» и скриншотов", kinozal_tabs(_kz) == [("7", 0, "Релиз")])
_kz_text = html_to_text(inline_tabs(_kz, [("Релиз", "Аудио 1: AC3 - Русский (MVO, LostFilm)<br>")]))
check("kinozal: вкладка «Релиз» попадает в текст",
      "Перевод: Профессиональный многоголосый" in _kz_text and "(MVO, LostFilm)" in _kz_text)
check("codec_segments: русская дорожка с типом, без иноязычной",
      codec_segments("Аудио 3: AC3, 2.0 ch - Русский (MVO, HDrezka Studio)\n"
                     "Аудио 6: E-AC3, 5.1 ch, 640 Kbps - English Dubbing\n"
                     "Перевод: Профессиональный") == ["AC3, 2.0 ch - Русский (MVO, HDrezka Studio)"])
check("same_studio: варианты написания",
      same_studio("DVD", "DVD Магия") and same_studio("Михаил Яроцкий", "М.Яроцкий/Kyberpunk")
      and not same_studio("LostFilm", "HDrezka"))
_page2 = [dict(_t("проф. многоголосый", None), studio_span=None, studio_confidence=None)]
merge_extra_tracks(_page2, [
    dict(_t("проф. многоголосый", "HDrezka"), studio_span="HDrezka Studio", studio_confidence=1),
    dict(_t("проф. многоголосый", "LostFilm"), studio_span="LostFilm", studio_confidence=1),
    dict(_t("одноголосый", None), studio_span=None, studio_confidence=None)],
    "codec", match=same_studio)
check("codec: студия вписана, новая добавлена, без студии — нет",
      [(x["studio"], x["source"]) for x in _page2]
      == [("HDrezka", "page+codec"), ("LostFilm", "codec")])

_dump = ("Меню\nГангстерленд / MobLand / 2025 / ПМ (LostFilm) / WEB\n"
         "Гангстерленд / 2026 / ДБ (X) / WEB-DL")
check("title_from_text: первая строка",
      "LostFilm" in title_from_text(_dump) and "ДБ" not in title_from_text(_dump))

from src.predict import refine_span  # noqa: E402

check("«Авторский, одноголосый - X» не режется",
      _split_entries("Авторский, одноголосый - Ю. Сербин") == ["Авторский, одноголосый - Ю. Сербин"])
check("перечисление голосностей режется",
      len(_split_entries("Одноголосый, многоголосый закадровый")) == 2)
_t4 = "Russian: 3/2 (L,C,R,l,r) + LFE ch |Дубляж"
check("span: раскладка каналов -> None", refine_span(_t4, (14, 23)) is None)
_t5 = "Русский MVO HDRezka Studio"
check("span: расширение до границ слова", _t5[slice(*refine_span(_t5, (12, 14)))] == "HDRezka")
check("span: один символ -> None", refine_span("М. Чадов", (0, 1)) is None)

import json as _json  # noqa: E402
import threading  # noqa: E402
import urllib.error  # noqa: E402
from http.server import BaseHTTPRequestHandler, HTTPServer  # noqa: E402

import src.fetch as _fetch  # noqa: E402


class _FakeFlare(BaseHTTPRequestHandler):
    def do_POST(self):
        req = _json.loads(self.rfile.read(int(self.headers["Content-Length"])))
        ok = req.get("cmd") == "request.get" and "rutracker.org" in req.get("url", "")
        body = _json.dumps({"status": "ok", "solution": {"response": "<html>Перевод: Дублированный | MovieDalen</html>"}}
                           if ok else {"status": "error", "message": "bad"}).encode()
        self.send_response(200)
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, *a):
        pass


_srv = HTTPServer(("127.0.0.1", 0), _FakeFlare)
threading.Thread(target=_srv.serve_forever, daemon=True).start()
_flare = f"http://127.0.0.1:{_srv.server_port}"


def _deny(url, timeout, cookie):
    raise urllib.error.HTTPError(url, 403, "Forbidden", {}, None)


_orig = _fetch._fetch_direct
_fetch._fetch_direct = _deny
try:
    h, tr = _fetch.fetch_html("https://rutracker.org/forum/viewtopic.php?t=1", flaresolverr=_flare)
    check("FlareSolverr: fallback при 403", "MovieDalen" in h and tr == "rutracker")
    try:
        _fetch.fetch_html("https://rutracker.org/forum/viewtopic.php?t=1")
        check("без FlareSolverr 403 пробрасывается", False)
    except urllib.error.HTTPError:
        check("без FlareSolverr 403 пробрасывается", True)
finally:
    _fetch._fetch_direct = _orig
    _srv.shutdown()


check("rutracker: форма логина в шапке не блокирует тему",
      blocked_reason('<html><a href="login.php"></a><input name="login_username"><div class="post_body">x</div></html>') is None)
check("rutracker: голая страница логина блокируется",
      blocked_reason('<html><a href="login.php"></a><input name="login_username"></html>') is not None)

_rt = ('<div class="post_body"><span class="post-b">Перевод 1</span>: Профессиональный '
       '(дублированный) <a href="#">Продубляж</a><br><span>Перевод 2</span>: Профессиональный '
       '(многоголосый закадровый) <a>AlexFilm</a><br>Субтитры: Русские</div>')
check("html: студия-ссылка внутри строки не рвёт дорожку",
      extract_segments(html_to_text(_rt)) == ["Профессиональный (дублированный) Продубляж",
                                              "Профессиональный (многоголосый закадровый) AlexFilm"])
check("«Дубляж, Продубляж» не режется",
      len(_split_entries("Russian AC3 |Дубляж, Продубляж|")) == 1)

_agg2 = aggregate([_t("проф. дубляж", "Продубляж"), _t("проф. дубляж", None),
                   _t("проф. двухголосый", None)])
check("aggregate: дорожка без студии того же типа схлопывается",
      [x["translation_type"] for x in _agg2["tracks"]] == ["проф. дубляж", "проф. двухголосый"])

print("== сегментатор (модель) ==")
try:
    from src.segmenter import LineSegmenter, is_test_release  # noqa: E402
    _seg = LineSegmenter("models/segmenter", "cpu")
    _fx = open("data/samples/rt-topic-6870900.html", encoding="utf-8").read()
    _lines = _seg.select_lines(html_to_text(_fx))
    check("сегментатор: найдены 4 строки «Перевод N» реальной страницы",
          sum(l.startswith("Перевод") for l in _lines) == 4, str(_lines[:5]))
    check("сегментатор: «Субтитры» и «Оригинальная аудиодорожка» не дорожки",
          not any(l.startswith(("Субтитры", "Оригинальная")) for l in _lines))
    check("разбиение страниц детерминировано", is_test_release("6870900") == is_test_release("6870900"))
except Exception as e:  # noqa: BLE001
    print(f"  пропущено (нет models/segmenter): {e}")

print(f"\nитог: {PASS} ok, {FAIL} failed")
sys.exit(1 if FAIL else 0)