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 | import argparse | 3 | import argparse |
| 4 | from datetime import UTC, datetime, timedelta | 4 | from datetime import UTC, datetime, timedelta |
| 5 | from pathlib import Path | 5 | from pathlib import Path |
| 6 | import re | ||
| 6 | 7 | ||
| 7 | from sqlalchemy import select | 8 | from sqlalchemy import select |
| 8 | 9 | ||
| @@ -11,19 +12,31 @@ from srht_contrib.db import make_session_factory | |||
| 11 | from srht_contrib.models import TrackedActor | 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 | def _iter_usernames(path: Path) -> list[str]: | 29 | def _iter_usernames(path: Path) -> list[str]: |
| 15 | usernames: list[str] = [] | 30 | usernames: list[str] = [] |
| 16 | seen: set[str] = set() | 31 | seen: set[str] = set() |
| 17 | for raw_line in path.read_text(encoding="utf-8").splitlines(): | 32 | for raw_line in path.read_text(encoding="utf-8").splitlines(): |
| 18 | username = raw_line.strip() | 33 | actor = _normalize_actor(raw_line) |
| 19 | if not username or username.startswith("#"): | 34 | if actor is None: |
| 20 | continue | 35 | continue |
| 21 | if not username.startswith("~"): | 36 | if actor in seen: |
| 22 | username = f"~{username}" | ||
| 23 | if username in seen: | ||
| 24 | continue | 37 | continue |
| 25 | seen.add(username) | 38 | seen.add(actor) |
| 26 | usernames.append(username) | 39 | usernames.append(actor) |
| 27 | return usernames | 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 | assert actors[0].next_poll_after == queued_at.replace(tzinfo=None) | 738 | assert actors[0].next_poll_after == queued_at.replace(tzinfo=None) |
| 739 | assert actors[1].next_poll_after == (queued_at + timedelta(seconds=60)).replace(tzinfo=None) | 739 | assert actors[1].next_poll_after == (queued_at + timedelta(seconds=60)).replace(tzinfo=None) |
| 740 | assert actors[2].next_poll_after == (queued_at + timedelta(seconds=120)).replace(tzinfo=None) | 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"] | ||