Commit daaeba8be8
Verified · cmc
Layout: unified · split
src/srht_contrib/scripts/enqueue_actors.py +20 −7
| @@ -3,6 +3,7 @@ from __future__ import annotations | ||
| 3 | 3 | import argparse |
| 4 | 4 | from datetime import UTC, datetime, timedelta |
| 5 | 5 | from pathlib import Path |
| 6 | import re | |
| 6 | 7 | |
| 7 | 8 | from sqlalchemy import select |
| 8 | 9 | |
| @@ -11,19 +12,31 @@ from srht_contrib.db import make_session_factory | ||
| 11 | 12 | from srht_contrib.models import TrackedActor |
| 12 | 13 | |
| 13 | 14 | |
| 15 | USERNAME_RE = re.compile(r"^[A-Za-z0-9](?:[A-Za-z0-9._-]{0,61}[A-Za-z0-9])?$") | |
| 16 | ||
| 17 | ||
| 18 | def _normalize_actor(raw_username: str) -> str | None: | |
| 19 | username = raw_username.strip() | |
| 20 | if not username or username.startswith("#"): | |
| 21 | return None | |
| 22 | if username.startswith("~"): | |
| 23 | username = username[1:] | |
| 24 | if not USERNAME_RE.fullmatch(username): | |
| 25 | return None | |
| 26 | return f"~{username}" | |
| 27 | ||
| 28 | ||
| 14 | 29 | def _iter_usernames(path: Path) -> list[str]: |
| 15 | 30 | usernames: list[str] = [] |
| 16 | 31 | seen: set[str] = set() |
| 17 | 32 | for raw_line in path.read_text(encoding="utf-8").splitlines(): |
| 18 | username = raw_line.strip() | |
| 19 | if not username or username.startswith("#"): | |
| 33 | actor = _normalize_actor(raw_line) | |
| 34 | if actor is None: | |
| 20 | 35 | continue |
| 21 | if not username.startswith("~"): | |
| 22 | username = f"~{username}" | |
| 23 | if username in seen: | |
| 36 | if actor in seen: | |
| 24 | 37 | continue |
| 25 | seen.add(username) | |
| 26 | usernames.append(username) | |
| 38 | seen.add(actor) | |
| 39 | usernames.append(actor) | |
| 27 | 40 | return usernames |
| 28 | 41 | |
| 29 | 42 | |
tests/test_ingestion.py +24
| @@ -738,3 +738,27 @@ def test_enqueue_actors_staggers_without_polling(tmp_path, monkeypatch) -> None: | ||
| 738 | 738 | assert actors[0].next_poll_after == queued_at.replace(tzinfo=None) |
| 739 | 739 | assert actors[1].next_poll_after == (queued_at + timedelta(seconds=60)).replace(tzinfo=None) |
| 740 | 740 | assert actors[2].next_poll_after == (queued_at + timedelta(seconds=120)).replace(tzinfo=None) |
| 741 | ||
| 742 | ||
| 743 | def test_enqueue_actors_skips_invalid_usernames(tmp_path, monkeypatch) -> None: | |
| 744 | database_path = tmp_path / "enqueue-invalid.db" | |
| 745 | username_path = tmp_path / "srht_usernames.txt" | |
| 746 | username_path.write_text("-0\n.\n~bad-\nvalid_user\nok.ok\n", encoding="utf-8") | |
| 747 | monkeypatch.setenv("DATABASE_URL", f"sqlite:///{database_path}") | |
| 748 | monkeypatch.setenv("SRHT_TOKEN", "test-token") | |
| 749 | monkeypatch.setenv("DEFAULT_ACTOR", "~ccleberg") | |
| 750 | ||
| 751 | from srht_contrib.db import Base, make_engine, make_session_factory | |
| 752 | ||
| 753 | settings = Settings() | |
| 754 | engine = make_engine(settings) | |
| 755 | Base.metadata.create_all(bind=engine) | |
| 756 | session_factory = make_session_factory(settings) | |
| 757 | ||
| 758 | inserted = enqueue_actors(Path(username_path), stagger_seconds=60) | |
| 759 | ||
| 760 | with session_factory() as db: | |
| 761 | actors = db.scalars(select(TrackedActor).order_by(TrackedActor.actor)).all() | |
| 762 | ||
| 763 | assert inserted == 2 | |
| 764 | assert [actor.actor for actor in actors] == ["~ok.ok", "~valid_user"] | |