287 lines
16 KiB
Python
287 lines
16 KiB
Python
"""Юнит-тесты 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")
|
||
|
||
_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()
|
||
|
||
|
||
print(f"\nитог: {PASS} ok, {FAIL} failed")
|
||
sys.exit(1 if FAIL else 0)
|