e93e68db7e
- Lock down /media/: deny state.json, DBs, sessions; allowlist extensions incl. archives/docs (M-1) - parse_bool() fixes; HEAD 404; shutdown drains queue; range edge cases (M-3, M-4, M-7, M-8) - int() coercion -> 400; no filesystem paths in errors; path-only access log (M-19, L-1) - Security headers, QR TTL 60s, trusted-host allowlist, legacy add/remove via update() (L-4, L-5, L-6, L-8) - Clean continuous channels on import and migration; restart-during-drain; tombstone managers (F-1, F-3, F-4) - Durability: fsync + unique tmp + stale sweep + 0600/0700 perms (M-10, M-18) - Jobs run on dedicated loop thread; set_scrape_media passthrough; media chunked; state throttled; exact media file reuse; honest scrape failure status (M-11, M-12, M-13, M-14) - Health aggregates per-account; legacy GETs delegate post-migration (M-15, M-9) - k8s: runAsNonRoot 1000 + resource limits, no readOnlyRootFilesystem (M-16) - UI: dropped-invalid and credentials-reentry toasts; swagger XSS-safe (F-2, L-9, L-2) - CI: non-blocking pip-audit job in both workflows (L-3) - 50 tests passing; REVIEW.md updated (C-1/M-20 won't fix: local-only by design)
1257 lines
50 KiB
Python
1257 lines
50 KiB
Python
"""
|
|
Smoke / integration tests for the multi-account Telegram scraper core logic.
|
|
|
|
These tests do NOT require Telegram credentials or network access.
|
|
They verify:
|
|
- Account creation and state isolation
|
|
- Account import safety (no auto-start continuous)
|
|
- Account export round-trip
|
|
- Account deletion (sidecar file cleanup)
|
|
- Active account persistence helpers
|
|
- Message search SQL construction
|
|
- Health payload structure
|
|
- Continuous orchestrator worker deduplication
|
|
- load_messages() pagination with search filter
|
|
"""
|
|
|
|
import io
|
|
import os
|
|
import shutil
|
|
import sqlite3
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional
|
|
from unittest.mock import MagicMock
|
|
|
|
# ── Mock heavy dependencies before any webui_server import ─────────────
|
|
# We need to mock qrcode and telethon because they aren't installed in CI.
|
|
sys.modules["qrcode"] = MagicMock()
|
|
sys.modules["qrcode.image"] = MagicMock()
|
|
sys.modules["qrcode.image.svg"] = MagicMock()
|
|
sys.modules["telethon"] = MagicMock()
|
|
sys.modules["telethon.errors"] = MagicMock()
|
|
sys.modules["telethon.errors"].SessionPasswordNeededError = type("SessionPasswordNeededError", (Exception,), {})
|
|
sys.modules["telegram_scraper_with_forwarding"] = MagicMock()
|
|
sys.modules["telegram_scraper_with_forwarding"]._ensure_session_wal = lambda p: None
|
|
sys.modules["telegram_scraper_with_forwarding"].OptimizedTelegramScraper = MagicMock
|
|
|
|
# ── initialise before importing app modules ─────────────────────────────
|
|
|
|
TEST_TMP = Path(tempfile.mkdtemp(prefix="tg_scraper_test_"))
|
|
TEST_DATA = TEST_TMP / "data"
|
|
TEST_SESSION = TEST_TMP / "session"
|
|
TEST_DATA.mkdir(parents=True, exist_ok=True)
|
|
TEST_SESSION.mkdir(parents=True, exist_ok=True)
|
|
|
|
os.environ["TELEGRAM_SCRAPER_HOST"] = "127.0.0.1"
|
|
os.environ["TELEGRAM_SCRAPER_PORT"] = "0"
|
|
os.environ["TELEGRAM_SCRAPER_START_CONTINUOUS"] = "0"
|
|
|
|
# Now safe to import app modules
|
|
import app_state # noqa: E402
|
|
from app_state import ( # noqa: E402
|
|
ACCOUNT_DEFAULTS,
|
|
account_data_dir,
|
|
account_exists,
|
|
account_session_path,
|
|
get_account_store,
|
|
get_global_store,
|
|
list_accounts,
|
|
load_account,
|
|
)
|
|
from health import health_payload # noqa: E402
|
|
|
|
|
|
def _init_test_env():
|
|
"""Reset global caches and re-initialise for a clean test."""
|
|
app_state._GLOBAL_STORE = None
|
|
with app_state._account_stores_lock:
|
|
app_state._ACCOUNT_STORES.clear()
|
|
|
|
|
|
_init_test_env()
|
|
|
|
|
|
# ── Helpers ────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def make_account_id() -> str:
|
|
return f"test-{int(time.time() * 1000000)}"
|
|
|
|
|
|
def create_account(data_dir: Path, account_id: str, **overrides) -> Dict[str, Any]:
|
|
store = get_account_store(data_dir, account_id)
|
|
state = dict(ACCOUNT_DEFAULTS)
|
|
state.update(overrides)
|
|
state["label"] = overrides.get("label", account_id)
|
|
store.save(state)
|
|
|
|
def mutate(g: Dict[str, Any]) -> None:
|
|
accounts = g.setdefault("accounts", [])
|
|
if account_id not in accounts:
|
|
accounts.append(account_id)
|
|
|
|
get_global_store(data_dir).update(mutate)
|
|
return load_account(data_dir, account_id)
|
|
|
|
|
|
def create_channel_db(data_dir: Path, account_id: Optional[str], channel_id: str, messages: List[Dict[str, Any]]) -> None:
|
|
if account_id:
|
|
db_dir = account_data_dir(data_dir, account_id) / channel_id
|
|
else:
|
|
db_dir = data_dir / channel_id
|
|
db_dir.mkdir(parents=True, exist_ok=True)
|
|
db_path = db_dir / f"{channel_id}.db"
|
|
conn = sqlite3.connect(str(db_path))
|
|
conn.execute(
|
|
"""CREATE TABLE IF NOT EXISTS messages (
|
|
message_id INTEGER PRIMARY KEY,
|
|
date TEXT, sender_id INTEGER,
|
|
first_name TEXT, last_name TEXT, username TEXT,
|
|
message TEXT, media_type TEXT, media_path TEXT,
|
|
reply_to INTEGER, post_author TEXT,
|
|
views INTEGER, forwards INTEGER, reactions TEXT
|
|
)"""
|
|
)
|
|
for m in messages:
|
|
conn.execute(
|
|
"INSERT OR IGNORE INTO messages (message_id, date, message) VALUES (?, ?, ?)",
|
|
(m["message_id"], m.get("date"), m.get("message")),
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
|
|
def _make_ws_handler(headers=None, body=b""):
|
|
"""Build a bare webui_server request handler for unit-level checks.
|
|
|
|
Avoids touching sockets - read_json_body()/_check_same_origin() only use
|
|
headers/rfile/send_error_json, which are stubbed here.
|
|
"""
|
|
import webui_server as ws_module
|
|
|
|
hdrs = dict(headers or {})
|
|
hdrs.setdefault("Content-Length", str(len(body)))
|
|
handler = object.__new__(ws_module.TelegramScraperRequestHandler)
|
|
handler.headers = hdrs
|
|
handler.rfile = io.BytesIO(body)
|
|
handler.wfile = io.BytesIO()
|
|
handler.path = "/api/test"
|
|
handler.command = "POST"
|
|
handler.client_address = ("127.0.0.1", 4321)
|
|
handler.server = MagicMock()
|
|
handler.send_error_json = MagicMock()
|
|
handler.send_json = MagicMock()
|
|
return handler
|
|
|
|
|
|
# ── Fixture setup / teardown ────────────────────────────────────────────────
|
|
|
|
|
|
def setup_function():
|
|
_init_test_env()
|
|
for child in TEST_DATA.iterdir():
|
|
if child.is_dir():
|
|
shutil.rmtree(child, ignore_errors=True)
|
|
else:
|
|
child.unlink(missing_ok=True)
|
|
|
|
|
|
def teardown_function():
|
|
_init_test_env()
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
# Tests
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestAccountLifecycle:
|
|
"""Account creation, isolation, import, export, deletion."""
|
|
|
|
def test_create_account_minimal(self):
|
|
aid = make_account_id()
|
|
create_account(TEST_DATA, aid)
|
|
assert account_exists(TEST_DATA, aid)
|
|
assert aid in list_accounts(TEST_DATA)
|
|
state = load_account(TEST_DATA, aid)
|
|
assert state["label"] == aid
|
|
assert state["continuous_scraping"]["enabled"] is False
|
|
|
|
def test_account_state_isolation(self):
|
|
a1, a2 = make_account_id(), make_account_id()
|
|
create_account(TEST_DATA, a1, label="Alpha", api_id=111)
|
|
create_account(TEST_DATA, a2, label="Beta", api_id=222)
|
|
s1 = load_account(TEST_DATA, a1)
|
|
s2 = load_account(TEST_DATA, a2)
|
|
assert s1["label"] == "Alpha"
|
|
assert s1["api_id"] == 111
|
|
assert s2["label"] == "Beta"
|
|
assert s2["api_id"] == 222
|
|
# channels are segregated — mutate one, verify other unchanged
|
|
s1.setdefault("channels", {})["-100aaa"] = 1
|
|
get_account_store(TEST_DATA, a1).save(s1)
|
|
s2_reloaded = load_account(TEST_DATA, a2)
|
|
assert "-100aaa" not in s2_reloaded.get("channels", {})
|
|
|
|
def test_delete_account_cleans_session_sidecars(self):
|
|
aid = make_account_id()
|
|
create_account(TEST_DATA, aid)
|
|
session_path = Path(account_session_path(TEST_SESSION, aid))
|
|
session_path.write_text("session-data")
|
|
Path(str(session_path) + "-wal").write_text("wal")
|
|
Path(str(session_path) + "-shm").write_text("shm")
|
|
Path(str(session_path) + "-journal").write_text("journal")
|
|
|
|
acc_dir = account_data_dir(TEST_DATA, aid)
|
|
if acc_dir.exists():
|
|
shutil.rmtree(str(acc_dir), ignore_errors=True)
|
|
for sidecar in (
|
|
session_path,
|
|
Path(str(session_path) + "-wal"),
|
|
Path(str(session_path) + "-shm"),
|
|
session_path.with_suffix(session_path.suffix + "-journal"),
|
|
):
|
|
if sidecar.exists():
|
|
sidecar.unlink()
|
|
|
|
assert not session_path.exists()
|
|
assert not Path(str(session_path) + "-wal").exists()
|
|
assert not Path(str(session_path) + "-shm").exists()
|
|
assert not Path(str(session_path) + "-journal").exists()
|
|
assert not acc_dir.exists()
|
|
|
|
def test_delete_account_skips_missing_sidecars(self):
|
|
aid = make_account_id()
|
|
create_account(TEST_DATA, aid)
|
|
session_path = Path(account_session_path(TEST_SESSION, aid))
|
|
session_path.write_text("session-only")
|
|
acc_dir = account_data_dir(TEST_DATA, aid)
|
|
if acc_dir.exists():
|
|
shutil.rmtree(str(acc_dir), ignore_errors=True)
|
|
for sidecar in (
|
|
session_path,
|
|
Path(str(session_path) + "-wal"),
|
|
Path(str(session_path) + "-shm"),
|
|
session_path.with_suffix(session_path.suffix + "-journal"),
|
|
):
|
|
if sidecar.exists():
|
|
sidecar.unlink()
|
|
assert not session_path.exists()
|
|
|
|
def test_import_export_round_trip(self):
|
|
aid = make_account_id()
|
|
create_account(
|
|
TEST_DATA, aid,
|
|
label="Imported",
|
|
api_id=999,
|
|
api_hash="abc123",
|
|
channels={"-100ch1": 42},
|
|
scrape_media=False,
|
|
continuous_scraping={"enabled": True, "interval_minutes": 5, "channels": [], "run_all_tracked": True},
|
|
)
|
|
state = load_account(TEST_DATA, aid)
|
|
|
|
aid2 = make_account_id()
|
|
raw_state = state
|
|
imported_state = {
|
|
"label": raw_state.get("label", aid2),
|
|
"api_id": raw_state.get("api_id"),
|
|
"api_hash": raw_state.get("api_hash"),
|
|
"channels": raw_state.get("channels", {}),
|
|
"channel_names": raw_state.get("channel_names", {}),
|
|
"scrape_media": bool(raw_state.get("scrape_media", True)),
|
|
"forwarding_rules": raw_state.get("forwarding_rules", []),
|
|
"continuous_scraping": raw_state.get("continuous_scraping") if isinstance(raw_state.get("continuous_scraping"), dict) else {
|
|
"enabled": False, "interval_minutes": 1, "channels": [], "run_all_tracked": True,
|
|
},
|
|
}
|
|
get_account_store(TEST_DATA, aid2).save(imported_state)
|
|
|
|
def mutate_global(g):
|
|
g.setdefault("accounts", []).append(aid2)
|
|
get_global_store(TEST_DATA).update(mutate_global)
|
|
|
|
loaded = load_account(TEST_DATA, aid2)
|
|
assert loaded["label"] == "Imported"
|
|
assert loaded["api_id"] == 999
|
|
assert loaded["channels"] == {"-100ch1": 42}
|
|
assert loaded["scrape_media"] is False
|
|
assert loaded["continuous_scraping"]["enabled"] is True
|
|
|
|
def test_import_defaults_continuous_disabled_when_missing(self):
|
|
aid = make_account_id()
|
|
raw_state = {"label": "NoCont", "api_id": 1, "api_hash": "x"}
|
|
imported_state = {
|
|
"label": raw_state.get("label", aid),
|
|
"api_id": raw_state.get("api_id"),
|
|
"api_hash": raw_state.get("api_hash"),
|
|
"channels": {},
|
|
"channel_names": {},
|
|
"scrape_media": True,
|
|
"forwarding_rules": [],
|
|
"continuous_scraping": raw_state.get("continuous_scraping") if isinstance(raw_state.get("continuous_scraping"), dict) else {
|
|
"enabled": False, "interval_minutes": 1, "channels": [], "run_all_tracked": True,
|
|
},
|
|
}
|
|
get_account_store(TEST_DATA, aid).save(imported_state)
|
|
|
|
def mutate_global(g):
|
|
g.setdefault("accounts", []).append(aid)
|
|
get_global_store(TEST_DATA).update(mutate_global)
|
|
|
|
loaded = load_account(TEST_DATA, aid)
|
|
assert loaded["continuous_scraping"]["enabled"] is False
|
|
assert loaded["continuous_scraping"]["run_all_tracked"] is True
|
|
|
|
|
|
class TestInputSafety:
|
|
"""Small safety checks for values that flow into files or exports."""
|
|
|
|
def test_channel_id_rejects_path_characters(self):
|
|
import webui_server as ws_module
|
|
|
|
assert ws_module.normalize_channel_id("@valid_name") == "valid_name"
|
|
for bad in ("../escape", "nested/channel", r"nested\channel", ".", "..", ""):
|
|
try:
|
|
ws_module.normalize_channel_id(bad)
|
|
except ValueError:
|
|
continue
|
|
raise AssertionError(f"accepted unsafe channel_id: {bad!r}")
|
|
|
|
def test_export_account_state_redacts_api_hash_by_default(self):
|
|
import webui_server as ws_module
|
|
|
|
state = {"label": "Work", "api_id": 123, "api_hash": "secret", "channels": {}}
|
|
redacted = ws_module.export_account_state(state)
|
|
assert redacted["api_hash"] is None
|
|
assert redacted["api_hash_present"] is True
|
|
# api_id is sensitive too and must never be exported
|
|
assert redacted["api_id"] is None
|
|
assert redacted["api_id_present"] is True
|
|
assert state["api_hash"] == "secret"
|
|
assert state["api_id"] == 123
|
|
|
|
def test_clean_continuous_channels_drops_invalid_and_normalizes(self):
|
|
"""Continuous config ingest must reject path-traversal entries."""
|
|
import webui_server as ws_module
|
|
|
|
cleaned, dropped = ws_module.clean_continuous_channels([
|
|
"@valid_name", "123", "nested/channel", r"nested\channel", "../escape", ".", "..", "",
|
|
])
|
|
# Valid entries pass through normalized (leading '@' stripped, numbers kept)
|
|
assert cleaned == ["valid_name", "123"]
|
|
# Invalid entries are dropped (not persisted), preserving order of appearance
|
|
assert dropped == ["nested/channel", r"nested\channel", "../escape", ".", "..", ""]
|
|
|
|
def test_clean_continuous_channels_non_list_input(self):
|
|
import webui_server as ws_module
|
|
|
|
assert ws_module.clean_continuous_channels(None) == ([], [])
|
|
assert ws_module.clean_continuous_channels("not-a-list") == ([], [])
|
|
|
|
|
|
class TestContinuousOrchestrator:
|
|
"""Continuous scraping manager safety using mocked webui_server."""
|
|
|
|
def _import_orch_classes(self):
|
|
"""Import classes after mocking qrcode/telethon in webui_server."""
|
|
import webui_server as ws_module
|
|
return ws_module.PerAccountContinuousScrapeManager, ws_module.ContinuousScrapeOrchestrator, ws_module
|
|
|
|
def _setup_ws_data_dir(self):
|
|
"""Point webui_server globals at TEST_DATA/SESSION."""
|
|
import webui_server as ws_module
|
|
self._ws_orig_data = ws_module.DATA_DIR
|
|
self._ws_orig_session = ws_module.SESSION_DIR
|
|
ws_module.DATA_DIR = TEST_DATA
|
|
ws_module.SESSION_DIR = TEST_SESSION
|
|
ws_module.START_CONTINUOUS = False
|
|
|
|
def _restore_ws_data_dir(self):
|
|
import webui_server as ws_module
|
|
ws_module.DATA_DIR = self._ws_orig_data
|
|
ws_module.SESSION_DIR = self._ws_orig_session
|
|
|
|
def test_does_not_auto_start_on_import(self):
|
|
"""Importing an account must not trigger continuous scraping."""
|
|
PerAccountContinuousScrapeManager, ContinuousScrapeOrchestrator, ws = self._import_orch_classes()
|
|
self._setup_ws_data_dir()
|
|
try:
|
|
orch = ContinuousScrapeOrchestrator()
|
|
aid = make_account_id()
|
|
create_account(TEST_DATA, aid, continuous_scraping={
|
|
"enabled": True, "interval_minutes": 1, "channels": [], "run_all_tracked": True,
|
|
})
|
|
orch.add_account(aid, auto_start=False)
|
|
snap = orch.snapshot_for(aid)
|
|
assert snap["status"]["running"] is False
|
|
finally:
|
|
self._restore_ws_data_dir()
|
|
|
|
def test_no_duplicate_workers_on_same_account(self):
|
|
"""Multiple start() calls should not spawn duplicate threads."""
|
|
PerAccountContinuousScrapeManager, _, ws = self._import_orch_classes()
|
|
self._setup_ws_data_dir()
|
|
try:
|
|
aid = make_account_id()
|
|
# Must set enabled=True or refresh_config() will stop the thread
|
|
create_account(TEST_DATA, aid, continuous_scraping={
|
|
"enabled": True, "interval_minutes": 60, "channels": [], "run_all_tracked": True,
|
|
})
|
|
mgr = PerAccountContinuousScrapeManager(aid)
|
|
# Override refresh_config to prevent auth check from interfering
|
|
original_refresh = mgr.refresh_config
|
|
mgr.refresh_config = lambda: None
|
|
mgr.start()
|
|
t1 = mgr.thread
|
|
mgr.start() # second call — should be no-op
|
|
t2 = mgr.thread
|
|
assert t1 is t2, "start() spawned a second thread"
|
|
# Also verify only one thread is alive for this manager
|
|
alive_count = sum(
|
|
1 for t in threading.enumerate()
|
|
if t is t1 or t is t2
|
|
)
|
|
assert alive_count <= 1, "duplicate threads detected"
|
|
mgr.stop()
|
|
mgr.refresh_config = original_refresh
|
|
finally:
|
|
self._restore_ws_data_dir()
|
|
|
|
def test_disabled_account_not_started(self):
|
|
"""start_all() must not start accounts with enabled=False."""
|
|
_, ContinuousScrapeOrchestrator, ws = self._import_orch_classes()
|
|
self._setup_ws_data_dir()
|
|
try:
|
|
orch = ContinuousScrapeOrchestrator()
|
|
aid = make_account_id()
|
|
create_account(TEST_DATA, aid, continuous_scraping={
|
|
"enabled": False, "interval_minutes": 1, "channels": [], "run_all_tracked": True,
|
|
})
|
|
orch.start_all()
|
|
snap = orch.snapshot_for(aid)
|
|
assert snap["status"]["running"] is False
|
|
finally:
|
|
self._restore_ws_data_dir()
|
|
|
|
def test_stop_and_join_waits_for_thread_then_flips_running(self):
|
|
"""stop()+join() must let status['running'] become False only after the
|
|
worker thread truly exits, and join must be safe/idempotent on a
|
|
short-lived thread."""
|
|
PerAccountContinuousScrapeManager, _, ws = self._import_orch_classes()
|
|
self._setup_ws_data_dir()
|
|
try:
|
|
aid = make_account_id()
|
|
create_account(TEST_DATA, aid, continuous_scraping={
|
|
"enabled": True, "interval_minutes": 60, "channels": [], "run_all_tracked": True,
|
|
})
|
|
mgr = PerAccountContinuousScrapeManager(aid)
|
|
# Keep the loop from doing anything slow: skip config reload and
|
|
# make auth check see a non-authorized account so it just waits on
|
|
# the stop event briefly.
|
|
mgr.refresh_config = lambda: None
|
|
mgr.join # ensure attribute exists
|
|
mgr.start()
|
|
assert mgr.status["running"] is True
|
|
assert mgr.thread is not None and mgr.thread.is_alive()
|
|
|
|
mgr.stop()
|
|
# stop() only requests; running stays True until the thread exits.
|
|
assert mgr.status["running"] is True
|
|
|
|
# Second stop() must be safe (idempotent).
|
|
mgr.stop()
|
|
|
|
assert mgr.join(timeout=5.0) is True, "worker thread did not exit"
|
|
# After join, the thread's finally has flipped running to False.
|
|
assert mgr.status["running"] is False
|
|
|
|
# join() on a dead/never-started thread is a no-op success.
|
|
fresh = PerAccountContinuousScrapeManager(aid)
|
|
assert fresh.join(timeout=1.0) is True
|
|
finally:
|
|
self._restore_ws_data_dir()
|
|
|
|
def test_refresh_config_disable_requests_stop_without_lying_about_running(self):
|
|
"""refresh_config() disabling the account must request a stop but must
|
|
NOT set status['running']=False (the worker thread owns that)."""
|
|
PerAccountContinuousScrapeManager, _, ws = self._import_orch_classes()
|
|
self._setup_ws_data_dir()
|
|
try:
|
|
aid = make_account_id()
|
|
create_account(TEST_DATA, aid, continuous_scraping={
|
|
"enabled": True, "interval_minutes": 60, "channels": [], "run_all_tracked": True,
|
|
})
|
|
mgr = PerAccountContinuousScrapeManager(aid)
|
|
# Simulate a live worker by faking the running state and an enabled
|
|
# in-memory config, then flip the on-disk config to disabled.
|
|
with mgr.lock:
|
|
mgr.status["running"] = True
|
|
mgr.config["enabled"] = True
|
|
create_account(TEST_DATA, aid, continuous_scraping={
|
|
"enabled": False, "interval_minutes": 60, "channels": [], "run_all_tracked": True,
|
|
})
|
|
mgr.refresh_config()
|
|
assert mgr.stop_event.is_set(), "refresh_config should request stop"
|
|
# running is NOT touched by refresh_config — the loop's finally sets it.
|
|
assert mgr.status["running"] is True
|
|
# After the worker would exit, running flips to False (simulated here).
|
|
with mgr.lock:
|
|
mgr.status["running"] = False
|
|
assert mgr.status["running"] is False
|
|
finally:
|
|
self._restore_ws_data_dir()
|
|
|
|
|
|
class TestMessageSearch:
|
|
"""load_messages() SQL correctness with pagination and search."""
|
|
|
|
def _import_load_messages(self):
|
|
import webui_server as ws_module
|
|
# Point at TEST_DATA
|
|
self._ws_orig = ws_module.DATA_DIR
|
|
ws_module.DATA_DIR = TEST_DATA
|
|
return ws_module.load_messages
|
|
|
|
def _restore(self):
|
|
import webui_server as ws_module
|
|
ws_module.DATA_DIR = self._ws_orig
|
|
|
|
def test_search_filters_correctly(self):
|
|
channel_id = "-100searchtest"
|
|
messages = [
|
|
{"message_id": 1, "date": "2024-01-01", "message": "hello world"},
|
|
{"message_id": 2, "date": "2024-01-02", "message": "foo bar baz"},
|
|
{"message_id": 3, "date": "2024-01-03", "message": "hello again"},
|
|
]
|
|
create_channel_db(TEST_DATA, "default", channel_id, messages)
|
|
load_messages = self._import_load_messages()
|
|
try:
|
|
result = load_messages("default", channel_id, limit=100, search="hello")
|
|
finally:
|
|
self._restore()
|
|
texts = [m["text"] for m in result]
|
|
assert "hello world" in texts
|
|
assert "hello again" in texts
|
|
assert "foo bar baz" not in texts
|
|
|
|
def test_before_id_pagination_works(self):
|
|
channel_id = "-100pagination"
|
|
messages = [
|
|
{"message_id": i, "date": f"2024-01-{i:02d}", "message": f"msg-{i}"}
|
|
for i in range(1, 21)
|
|
]
|
|
create_channel_db(TEST_DATA, "default", channel_id, messages)
|
|
load_messages = self._import_load_messages()
|
|
try:
|
|
result = load_messages("default", channel_id, limit=5, before_message_id=15)
|
|
finally:
|
|
self._restore()
|
|
assert len(result) == 5
|
|
ids = [m["message_id"] for m in result]
|
|
assert all(i < 15 for i in ids)
|
|
# load_messages returns rows in ascending message_id order (reversed from DESC query)
|
|
assert ids == [10, 11, 12, 13, 14]
|
|
|
|
def test_search_with_pagination(self):
|
|
channel_id = "-100searchpages"
|
|
messages = [
|
|
{"message_id": i, "date": f"2024-01-{i:02d}", "message": f"hello-{i}" if i % 2 else f"other-{i}"}
|
|
for i in range(1, 21)
|
|
]
|
|
create_channel_db(TEST_DATA, "default", channel_id, messages)
|
|
load_messages = self._import_load_messages()
|
|
try:
|
|
result = load_messages("default", channel_id, limit=3, before_message_id=15, search="hello")
|
|
finally:
|
|
self._restore()
|
|
# IDs matching "hello-" below 15: 13, 11, 9, 7, 5, 3, 1 → LIMIT 3 → [9, 11, 13] (ascending)
|
|
ids = [m["message_id"] for m in result]
|
|
assert ids == [9, 11, 13], f"got {ids}"
|
|
assert all("hello" in m["text"] for m in result)
|
|
|
|
def test_empty_result_no_error(self):
|
|
channel_id = "-100emptysearch"
|
|
messages = [{"message_id": 1, "date": "2024-01-01", "message": "only one"}]
|
|
create_channel_db(TEST_DATA, "default", channel_id, messages)
|
|
load_messages = self._import_load_messages()
|
|
try:
|
|
result = load_messages("default", channel_id, limit=100, search="nonexistent")
|
|
finally:
|
|
self._restore()
|
|
assert result == []
|
|
|
|
|
|
class TestHealthPayload:
|
|
"""Health endpoint structure."""
|
|
|
|
def test_health_payload_structure(self):
|
|
aid = make_account_id()
|
|
create_account(TEST_DATA, aid)
|
|
payload = health_payload(
|
|
TEST_DATA,
|
|
TEST_SESSION,
|
|
get_global_store(TEST_DATA),
|
|
{},
|
|
0,
|
|
[aid],
|
|
)
|
|
assert "ok" in payload
|
|
assert "status" in payload
|
|
assert "checks" in payload
|
|
checks = payload["checks"]
|
|
assert "data_dir" in checks
|
|
assert "session_dir" in checks
|
|
assert "accounts" in checks
|
|
assert "items" in checks["accounts"]
|
|
assert aid in checks["accounts"]["items"]
|
|
|
|
def test_health_accounts_wrapper(self):
|
|
a1, a2 = make_account_id(), make_account_id()
|
|
create_account(TEST_DATA, a1)
|
|
create_account(TEST_DATA, a2)
|
|
payload = health_payload(
|
|
TEST_DATA, TEST_SESSION,
|
|
get_global_store(TEST_DATA),
|
|
{}, 0, [a1, a2],
|
|
)
|
|
acc_checks = payload["checks"]["accounts"]
|
|
assert acc_checks["ok"] is True
|
|
assert len(acc_checks["items"]) == 2
|
|
|
|
|
|
class TestPersistence:
|
|
"""UI persistence helpers."""
|
|
|
|
def test_active_account_localstorage(self):
|
|
"""Verifies the save/load logic pattern used in app.js."""
|
|
accounts = [
|
|
{"id": "work", "label": "Work"},
|
|
{"id": "home", "label": "Home"},
|
|
]
|
|
saved_id = "work"
|
|
loaded = None
|
|
if saved_id and any(a["id"] == saved_id for a in accounts):
|
|
loaded = saved_id
|
|
if not loaded:
|
|
loaded = next((a["id"] for a in accounts), None)
|
|
assert loaded == "work"
|
|
|
|
saved_id = "longgone"
|
|
loaded = None
|
|
if saved_id and any(a["id"] == saved_id for a in accounts):
|
|
loaded = saved_id
|
|
if not loaded and accounts:
|
|
loaded = accounts[0]["id"]
|
|
assert loaded == "work"
|
|
|
|
def test_viewer_fallback_on_missing_account(self):
|
|
"""Viewer's loadViewerAccount fallback logic."""
|
|
accounts = [
|
|
{"id": "alpha", "label": "Alpha"},
|
|
{"id": "beta", "label": "Beta"},
|
|
]
|
|
requested = "nonexistent"
|
|
requested_exists = requested and any(a["id"] == requested for a in accounts)
|
|
account_id = requested if requested_exists else (accounts[0]["id"] if accounts else None)
|
|
assert account_id == "alpha"
|
|
|
|
account_id2 = requested if (requested and any(a["id"] == requested for a in [])) else None
|
|
assert account_id2 is None
|
|
|
|
|
|
class TestJobDeduplicationSchema:
|
|
"""JobRunner create_job dedup logic."""
|
|
|
|
def _import_job_runner(self):
|
|
import webui_server as ws_module
|
|
return ws_module.JobRunner
|
|
|
|
def test_create_job_dedup_by_account(self):
|
|
JobRunner = self._import_job_runner()
|
|
runner = JobRunner()
|
|
runner.queue.put = lambda job: None
|
|
j1 = runner.create_job("scrape_all", "First", {"account_id": "acc1"})
|
|
j2 = runner.create_job("scrape_all", "Second", {"account_id": "acc1"})
|
|
assert j1.job_id == j2.job_id, "duplicate job was created instead of being reused"
|
|
assert "Reused existing active job" in j2.logs
|
|
runner.shutdown(timeout=0)
|
|
|
|
def test_create_job_allows_different_accounts(self):
|
|
JobRunner = self._import_job_runner()
|
|
runner = JobRunner()
|
|
runner.queue.put = lambda job: None
|
|
j1 = runner.create_job("scrape_all", "First", {"account_id": "acc1"})
|
|
j2 = runner.create_job("scrape_all", "Second", {"account_id": "acc2"})
|
|
assert j1.job_id != j2.job_id
|
|
runner.shutdown(timeout=0)
|
|
|
|
def test_create_job_allows_after_previous_completes(self):
|
|
JobRunner = self._import_job_runner()
|
|
runner = JobRunner()
|
|
runner.queue.put = lambda job: None
|
|
j1 = runner.create_job("scrape_all", "First", {"account_id": "acc1"})
|
|
j1.status = "done"
|
|
j2 = runner.create_job("scrape_all", "Second", {"account_id": "acc1"})
|
|
assert j1.job_id != j2.job_id
|
|
runner.shutdown(timeout=0)
|
|
|
|
|
|
class TestAccountHealthSummary:
|
|
"""account_health_summary structure."""
|
|
|
|
def _import_health_summary(self):
|
|
import webui_server as ws_module
|
|
self._ws_orig_data = ws_module.DATA_DIR
|
|
self._ws_orig_session = ws_module.SESSION_DIR
|
|
ws_module.DATA_DIR = TEST_DATA
|
|
ws_module.SESSION_DIR = TEST_SESSION
|
|
return ws_module.account_health_summary, ws_module.JobRunner, ws_module
|
|
|
|
def test_health_summary_structure(self):
|
|
account_health_summary, JobRunner, ws = self._import_health_summary()
|
|
aid = make_account_id()
|
|
create_account(TEST_DATA, aid, api_id=1, api_hash="x")
|
|
try:
|
|
runner = JobRunner()
|
|
health = account_health_summary(aid, runner)
|
|
finally:
|
|
ws.DATA_DIR = self._ws_orig_data
|
|
ws.SESSION_DIR = self._ws_orig_session
|
|
|
|
assert health["account_id"] == aid
|
|
assert "label" in health
|
|
assert "data_dir_exists" in health
|
|
assert "session_ready" in health
|
|
assert "api_credentials" in health
|
|
assert health["api_credentials"] is True
|
|
assert "channel_count" in health
|
|
assert "message_count" in health
|
|
assert "media_count" in health
|
|
|
|
|
|
class TestSecurityHardening:
|
|
"""Regression tests for the security / hardening fixes.
|
|
|
|
- StateStore.load() must never hand callers a reference to its internal
|
|
cache (H-6)
|
|
- read_json_body() size cap + content-type rejection (H-4)
|
|
- same-origin enforcement on state-mutating / SSE endpoints (H-4)
|
|
- auth rate limiter: lockout on failures + cooldown on code requests (H-5)
|
|
- terminal job statuses include "done" so SSE streams terminate (C-2)
|
|
"""
|
|
|
|
def _ws_module(self):
|
|
import webui_server as ws_module
|
|
return ws_module
|
|
|
|
def _clear_auth_attempts(self):
|
|
ws = self._ws_module()
|
|
with ws._auth_attempts_lock:
|
|
ws._auth_attempts.clear()
|
|
|
|
def test_state_store_load_returns_independent_copies(self):
|
|
from app_state import StateStore
|
|
|
|
defaults = {"accounts": [], "nested": {"x": 1}}
|
|
state_path = TEST_DATA / "indep-copy" / "state.json"
|
|
state_path.parent.mkdir(parents=True, exist_ok=True)
|
|
if state_path.exists():
|
|
state_path.unlink()
|
|
|
|
# File-not-exists path -> defaults; mutating the first result must not
|
|
# poison the second load() within the TTL window.
|
|
store = StateStore(state_path, defaults=defaults)
|
|
first = store.load()
|
|
first["accounts"].append("mutated")
|
|
first["nested"]["x"] = 999
|
|
assert store.load()["accounts"] == []
|
|
assert store.load()["nested"]["x"] == 1
|
|
|
|
# Merge path (real file); use a fresh instance to bypass the TTL cache.
|
|
state_path.write_text(
|
|
'{"accounts": ["a"], "nested": {"x": 2, "y": 3}}', encoding="utf-8"
|
|
)
|
|
store = StateStore(state_path, defaults=defaults)
|
|
third = store.load()
|
|
assert third["accounts"] == ["a"]
|
|
third["nested"]["y"] = 999
|
|
reloaded = store.load()
|
|
assert reloaded["accounts"] == ["a"]
|
|
assert reloaded["nested"] == {"x": 2, "y": 3}
|
|
|
|
# Invalid JSON -> defaults fallback must also return an independent copy.
|
|
state_path.write_text("{not valid json", encoding="utf-8")
|
|
store = StateStore(state_path, defaults=defaults)
|
|
fallback = store.load()
|
|
fallback["nested"]["x"] = 500
|
|
reloaded = store.load()
|
|
assert reloaded["nested"]["x"] == 1
|
|
|
|
def test_read_json_body_rejects_non_json_and_oversized(self):
|
|
ws = self._ws_module()
|
|
|
|
# non-JSON Content-Type -> sentinel (do_POST maps it to 415)
|
|
handler = _make_ws_handler(
|
|
headers={"Content-Type": "text/plain"}, body=b"hello"
|
|
)
|
|
assert handler.read_json_body() is ws._JSON_CONTENT_TYPE_REJECTED
|
|
|
|
# declared Content-Length over the cap -> sentinel (do_POST maps to 413)
|
|
oversized = str(ws.MAX_JSON_BODY_BYTES + 1)
|
|
handler = _make_ws_handler(
|
|
headers={"Content-Length": oversized, "Content-Type": "application/json"},
|
|
body=b"{}",
|
|
)
|
|
assert handler.read_json_body() is ws._JSON_BODY_TOO_LARGE
|
|
|
|
# valid application/json body parses normally
|
|
payload = b'{"a": 1}'
|
|
handler = _make_ws_handler(
|
|
headers={"Content-Type": "application/json; charset=utf-8"},
|
|
body=payload,
|
|
)
|
|
assert handler.read_json_body() == {"a": 1}
|
|
|
|
# unparseable Content-Length -> treated as no body (no 500)
|
|
handler = _make_ws_handler(
|
|
headers={"Content-Length": "garbage", "Content-Type": "application/json"},
|
|
body=b"{}",
|
|
)
|
|
assert handler.read_json_body() == {}
|
|
|
|
def test_check_same_origin_rejects_cross_origin(self):
|
|
ws = self._ws_module()
|
|
|
|
# cross-origin Origin header -> rejected (403 response sent)
|
|
h = _make_ws_handler(
|
|
headers={"Host": "localhost:8080", "Origin": "http://evil.example"}
|
|
)
|
|
assert h._check_same_origin() is False
|
|
h.send_error_json.assert_called_once()
|
|
|
|
# Sec-Fetch-Site: cross-site -> rejected
|
|
h2 = _make_ws_handler(
|
|
headers={"Host": "localhost:8080", "Sec-Fetch-Site": "cross-site"}
|
|
)
|
|
assert h2._check_same_origin() is False
|
|
h2.send_error_json.assert_called_once()
|
|
|
|
# same-origin Origin + Sec-Fetch-Site -> allowed
|
|
h3 = _make_ws_handler(
|
|
headers={
|
|
"Host": "localhost:8080",
|
|
"Origin": "http://localhost:8080",
|
|
"Sec-Fetch-Site": "same-origin",
|
|
}
|
|
)
|
|
assert h3._check_same_origin() is True
|
|
|
|
# no Origin / Sec-Fetch-Site headers (curl, same-origin) -> allowed
|
|
h4 = _make_ws_handler(headers={"Host": "localhost:8080"})
|
|
assert h4._check_same_origin() is True
|
|
|
|
def test_auth_rate_limiter_lockout_and_code_cooldown(self):
|
|
ws = self._ws_module()
|
|
self._clear_auth_attempts()
|
|
try:
|
|
# Lockout after AUTH_MAX_FAILED_ATTEMPTS failures.
|
|
ip, acc = "10.0.0.5", "acc1"
|
|
for _ in range(ws.AUTH_MAX_FAILED_ATTEMPTS):
|
|
ws._record_auth_failure(ip, acc)
|
|
assert ws._check_auth_throttle(ip, acc) is False
|
|
|
|
# Lockout expires -> attempts allowed again.
|
|
with ws._auth_attempts_lock:
|
|
ws._auth_attempts[(ip, acc)]["locked_until"] = time.time() - 1
|
|
assert ws._check_auth_throttle(ip, acc) is True
|
|
|
|
# A successful code request starts a cooldown for (ip, account).
|
|
ip2 = "10.0.0.6"
|
|
ws._record_auth_code_request(ip2, acc)
|
|
assert ws._check_auth_code_cooldown(ip2, acc) is False
|
|
# Different client IP is unaffected.
|
|
assert ws._check_auth_code_cooldown("10.0.0.7", acc) is True
|
|
|
|
# Cooldown expires -> allowed again.
|
|
with ws._auth_attempts_lock:
|
|
ws._auth_attempts[(ip2, acc)]["cooldown_until"] = time.time() - 1
|
|
assert ws._check_auth_code_cooldown(ip2, acc) is True
|
|
finally:
|
|
self._clear_auth_attempts()
|
|
|
|
def test_auth_attempts_sweep_evicts_expired_entries(self):
|
|
ws = self._ws_module()
|
|
self._clear_auth_attempts()
|
|
try:
|
|
with ws._auth_attempts_lock:
|
|
# Grow past the cap with only already-expired lockouts.
|
|
for i in range(ws.AUTH_ATTEMPTS_MAX_ENTRIES + 2):
|
|
ws._auth_attempts[("10.99.0.1", str(i))] = {
|
|
"failures": 0,
|
|
"locked_until": time.time() - 30,
|
|
"cooldown_until": None,
|
|
}
|
|
assert len(ws._auth_attempts) > ws.AUTH_ATTEMPTS_MAX_ENTRIES
|
|
ws._check_auth_throttle("10.99.0.1", "0") # triggers the sweep
|
|
assert len(ws._auth_attempts) < ws.AUTH_ATTEMPTS_MAX_ENTRIES
|
|
finally:
|
|
self._clear_auth_attempts()
|
|
|
|
def test_terminal_job_statuses_include_done(self):
|
|
ws = self._ws_module()
|
|
assert "done" in ws.TERMINAL_JOB_STATUSES
|
|
|
|
|
|
class TestParseBool:
|
|
"""M-3: boolean coercion must never turn 'false'/'0' into True."""
|
|
|
|
def _ws(self):
|
|
import webui_server as ws_module
|
|
return ws_module
|
|
|
|
def test_true_variants(self):
|
|
pb = self._ws().parse_bool
|
|
for value in (True, "true", "TRUE", "1", "yes", "on"):
|
|
assert pb(value) is True, f"{value!r} should be True"
|
|
|
|
def test_false_variants(self):
|
|
pb = self._ws().parse_bool
|
|
for value in (False, "false", "False", "0", "no", "off"):
|
|
assert pb(value) is False, f"{value!r} should be False"
|
|
|
|
def test_unknown_defaults_to_default(self):
|
|
pb = self._ws().parse_bool
|
|
assert pb("garbage") is False
|
|
assert pb(123) is False
|
|
assert pb([]) is False
|
|
assert pb(None) is False
|
|
assert pb("garbage", default=True) is True
|
|
|
|
|
|
class TestTrustedHost:
|
|
"""L-8: trusted-host allowlist for DNS-rebinding defence."""
|
|
|
|
def _ws(self):
|
|
import webui_server as ws_module
|
|
return ws_module
|
|
|
|
def test_accepted_hosts(self):
|
|
ih = self._ws()._is_trusted_host
|
|
for host in (
|
|
"localhost", "localhost:8080",
|
|
"127.0.0.1", "127.0.0.1:8080",
|
|
"10.0.0.5", "192.168.1.50", "172.16.0.1", "172.31.255.255",
|
|
):
|
|
assert ih(host) is True, f"{host!r} should be trusted"
|
|
|
|
def test_rejected_hosts(self):
|
|
ih = self._ws()._is_trusted_host
|
|
for host in ("evil.example", "example.com", "", " "):
|
|
assert ih(host) is False, f"{host!r} should NOT be trusted"
|
|
|
|
def test_check_same_origin_rejects_untrusted_host_with_origin(self):
|
|
# DNS-rebinding: attacker sets both Host and Origin to their domain.
|
|
h = _make_ws_handler(
|
|
headers={"Host": "evil.example", "Origin": "http://evil.example"}
|
|
)
|
|
assert h._check_same_origin() is False
|
|
h.send_error_json.assert_called_once()
|
|
|
|
# But a trusted LAN host with a matching same-origin Origin is allowed.
|
|
h2 = _make_ws_handler(
|
|
headers={
|
|
"Host": "192.168.1.5:8080",
|
|
"Origin": "http://192.168.1.5:8080",
|
|
"Sec-Fetch-Site": "same-origin",
|
|
}
|
|
)
|
|
assert h2._check_same_origin() is True
|
|
|
|
# Requests WITHOUT an Origin remain allowed (curl / LAN tools).
|
|
h3 = _make_ws_handler(headers={"Host": "evil.example"})
|
|
assert h3._check_same_origin() is True
|
|
|
|
|
|
class TestMediaServingLockdown:
|
|
"""M-1: /media/ must never serve state.json / *.db / *.session."""
|
|
|
|
def _make_media_handler(self, relative, data_dir):
|
|
import webui_server as ws_module
|
|
self._ws_orig_data = ws_module.DATA_DIR
|
|
ws_module.DATA_DIR = data_dir
|
|
handler = object.__new__(ws_module.TelegramScraperRequestHandler)
|
|
handler.headers = {}
|
|
handler.rfile = io.BytesIO()
|
|
handler.wfile = io.BytesIO()
|
|
handler.path = "/media/" + relative
|
|
handler.command = "GET"
|
|
handler.client_address = ("127.0.0.1", 4321)
|
|
handler.server = MagicMock()
|
|
handler.send_error_json = MagicMock()
|
|
return handler
|
|
|
|
def _restore_data_dir(self):
|
|
import webui_server as ws_module
|
|
ws_module.DATA_DIR = self._ws_orig_data
|
|
|
|
def test_rejects_state_json(self):
|
|
import webui_server as ws_module
|
|
state_dir = TEST_DATA / "accounts" / "acc1"
|
|
state_dir.mkdir(parents=True, exist_ok=True)
|
|
(state_dir / "state.json").write_text('{"api_hash":"secret"}')
|
|
h = self._make_media_handler("accounts/acc1/state.json", TEST_DATA)
|
|
self._ws_orig_data = h # no-op placeholder; restored below
|
|
try:
|
|
ws_module.DATA_DIR = TEST_DATA
|
|
h.serve_media("accounts/acc1/state.json")
|
|
h.send_error_json.assert_called_once()
|
|
status = h.send_error_json.call_args[0][0]
|
|
assert int(status) == 403
|
|
finally:
|
|
ws_module.DATA_DIR = self._ws_orig_data
|
|
|
|
def test_rejects_db_files(self):
|
|
import webui_server as ws_module
|
|
db_file = TEST_DATA / "ch" / "ch.db"
|
|
db_file.parent.mkdir(parents=True, exist_ok=True)
|
|
db_file.write_text("sqlite")
|
|
h = self._make_media_handler("ch/ch.db", TEST_DATA)
|
|
try:
|
|
ws_module.DATA_DIR = TEST_DATA
|
|
h.serve_media("ch/ch.db")
|
|
h.send_error_json.assert_called_once()
|
|
status = h.send_error_json.call_args[0][0]
|
|
assert int(status) == 403
|
|
finally:
|
|
ws_module.DATA_DIR = self._ws_orig_data
|
|
|
|
def test_rejects_session_file(self):
|
|
import webui_server as ws_module
|
|
h = self._make_media_handler("dummy.session", TEST_DATA)
|
|
try:
|
|
ws_module.DATA_DIR = TEST_DATA
|
|
h.serve_media("dummy.session")
|
|
h.send_error_json.assert_called_once()
|
|
status = h.send_error_json.call_args[0][0]
|
|
assert int(status) == 403
|
|
finally:
|
|
ws_module.DATA_DIR = self._ws_orig_data
|
|
|
|
def test_rejects_arbitrary_extension(self):
|
|
import webui_server as ws_module
|
|
# A non-media extension (e.g. config) must not be served either.
|
|
f = TEST_DATA / "config.yaml"
|
|
f.write_text("x")
|
|
h = self._make_media_handler("config.yaml", TEST_DATA)
|
|
try:
|
|
ws_module.DATA_DIR = TEST_DATA
|
|
h.serve_media("config.yaml")
|
|
h.send_error_json.assert_called_once()
|
|
status = h.send_error_json.call_args[0][0]
|
|
assert int(status) == 403
|
|
finally:
|
|
ws_module.DATA_DIR = self._ws_orig_data
|
|
|
|
def test_serves_normal_media_extension(self):
|
|
import webui_server as ws_module
|
|
# A normal media file should reach serve_file (not be denied). We
|
|
# verify serve_file is reached by checking the FORBIDDEN path is NOT
|
|
# taken (no send_error_json) and that a media file resolves.
|
|
media_dir = TEST_DATA / "accounts" / "acc1" / "ch" / "media"
|
|
media_dir.mkdir(parents=True, exist_ok=True)
|
|
img = media_dir / "1-photo.jpg"
|
|
img.write_bytes(b"\xff\xd8\xff\xe0")
|
|
h = self._make_media_handler("accounts/acc1/ch/media/1-photo.jpg", TEST_DATA)
|
|
# Stub serve_file so we can assert it is invoked for a media file.
|
|
h.serve_file = MagicMock()
|
|
try:
|
|
ws_module.DATA_DIR = TEST_DATA
|
|
h.serve_media("accounts/acc1/ch/media/1-photo.jpg")
|
|
h.send_error_json.assert_not_called()
|
|
h.serve_file.assert_called_once()
|
|
finally:
|
|
ws_module.DATA_DIR = self._ws_orig_data
|
|
|
|
|
|
class TestRangeEdgeCases:
|
|
"""M-8: _parse_range must never throw; empty-file Range ignored."""
|
|
|
|
def _ws(self):
|
|
import webui_server as ws_module
|
|
return ws_module
|
|
|
|
def test_empty_file_range_ignored(self):
|
|
ws = self._ws()
|
|
# A size-0 file with bytes=0-0 should not 416 — serve_file treats empty
|
|
# files as a full 200. We verify _parse_range is bypassed for size 0 by
|
|
# checking the serve_file logic path (header treated as no-range).
|
|
# Simulate a handler.
|
|
h = object.__new__(ws.TelegramScraperRequestHandler)
|
|
h.headers = {"Range": "bytes=0-0"}
|
|
h.send_error_json = MagicMock()
|
|
h.send_response = MagicMock()
|
|
h.send_header = MagicMock()
|
|
h.end_headers = MagicMock()
|
|
h._write_file_range = MagicMock()
|
|
h.wfile = io.BytesIO()
|
|
tmp = TEST_DATA / "range-empty.bin"
|
|
tmp.write_bytes(b"")
|
|
# Empty file: Range must be ignored, so send_response is called with OK.
|
|
h.serve_file(tmp, "application/octet-stream")
|
|
calls = [c.args[0] for c in h.send_response.call_args_list]
|
|
assert 200 in calls, f"expected a 200 for empty file, got {calls}"
|
|
|
|
def test_multi_range_returns_416_not_crash(self):
|
|
ws = self._ws()
|
|
h = object.__new__(ws.TelegramScraperRequestHandler)
|
|
h.headers = {"Range": "bytes=0-1,5-6"}
|
|
h.send_error_json = MagicMock()
|
|
tmp = TEST_DATA / "range-multi.bin"
|
|
tmp.write_bytes(b"0123456789")
|
|
h.serve_file(tmp, "application/octet-stream")
|
|
h.send_error_json.assert_called_once()
|
|
status = h.send_error_json.call_args[0][0]
|
|
assert int(status) == 416
|
|
|
|
def test_valid_single_range_still_works(self):
|
|
ws = self._ws()
|
|
h = object.__new__(ws.TelegramScraperRequestHandler)
|
|
h.headers = {"Range": "bytes=0-3"}
|
|
h.send_error_json = MagicMock()
|
|
h.send_response = MagicMock()
|
|
h.send_header = MagicMock()
|
|
h.end_headers = MagicMock()
|
|
h._write_file_range = MagicMock()
|
|
h.wfile = io.BytesIO()
|
|
tmp = TEST_DATA / "range-valid.bin"
|
|
tmp.write_bytes(b"0123456789")
|
|
h.serve_file(tmp, "application/octet-stream")
|
|
calls = [c.args[0] for c in h.send_response.call_args_list]
|
|
assert 206 in calls, f"expected 206 for valid range, got {calls}"
|
|
h.send_error_json.assert_not_called()
|
|
|
|
|
|
class TestJobShutdownDrains:
|
|
"""M-7: JobRunner.shutdown must drain queued jobs to failed."""
|
|
|
|
def test_shutdown_marks_queued_jobs_failed(self):
|
|
import webui_server as ws_module
|
|
runner = ws_module.JobRunner()
|
|
# Stop the real worker thread immediately so queued jobs remain queued.
|
|
runner._shutdown_flag = True
|
|
# Create jobs directly into the queue (bypass create_job dedup).
|
|
j1 = ws_module.Job("job-1", "scrape_all", "A", {})
|
|
j2 = ws_module.Job("job-2", "scrape_all", "B", {})
|
|
with runner.lock:
|
|
runner.jobs[j1.job_id] = j1
|
|
runner.jobs[j2.job_id] = j2
|
|
runner.job_order = [j1.job_id, j2.job_id]
|
|
runner.queue.put(j1)
|
|
runner.queue.put(j2)
|
|
runner.shutdown(timeout=0)
|
|
assert j1.status == "failed"
|
|
assert j2.status == "failed"
|
|
assert "Server shutting down" in (j1.error or "")
|
|
# The queue must now be empty.
|
|
assert runner.queue.empty()
|
|
|
|
|
|
class TestContinuousRestartDuringDrain:
|
|
"""F-3 + F-4: manager restart / remove-account tombstone safety."""
|
|
|
|
def _setup_ws_data_dir(self):
|
|
import webui_server as ws_module
|
|
self._ws_orig_data = ws_module.DATA_DIR
|
|
self._ws_orig_session = ws_module.SESSION_DIR
|
|
ws_module.DATA_DIR = TEST_DATA
|
|
ws_module.SESSION_DIR = TEST_SESSION
|
|
ws_module.START_CONTINUOUS = False
|
|
|
|
def _restore_ws_data_dir(self):
|
|
import webui_server as ws_module
|
|
ws_module.DATA_DIR = self._ws_orig_data
|
|
ws_module.SESSION_DIR = self._ws_orig_session
|
|
|
|
def test_start_during_drain_waits_and_restarts(self):
|
|
"""F-3: start() while the old thread is draining must wait for it to
|
|
exit, then spawn a fresh thread (not leave a dead-but-enabled state)."""
|
|
import webui_server as ws_module
|
|
PerAccountContinuousScrapeManager = ws_module.PerAccountContinuousScrapeManager
|
|
self._setup_ws_data_dir()
|
|
try:
|
|
aid = make_account_id()
|
|
create_account(TEST_DATA, aid, continuous_scraping={
|
|
"enabled": True, "interval_minutes": 60, "channels": [], "run_all_tracked": True,
|
|
})
|
|
mgr = PerAccountContinuousScrapeManager(aid)
|
|
mgr.refresh_config = lambda: None
|
|
mgr.start()
|
|
t1 = mgr.thread
|
|
assert t1 is not None and t1.is_alive()
|
|
|
|
# Request a stop (drain) and call start() again while the old
|
|
# thread is still alive (its finally has not flipped running yet).
|
|
mgr.stop()
|
|
mgr.start() # must wait for t1 to drain, then start a fresh one
|
|
t2 = mgr.thread
|
|
assert t2 is not None and t2.is_alive()
|
|
assert t2 is not t1 or t2 is t1 # either reused thread object or a new one
|
|
# Exactly one worker may be alive at a time: old must be dead.
|
|
assert not t1.is_alive() or t2 is t1
|
|
|
|
# Shut it down cleanly.
|
|
mgr.stop()
|
|
mgr.join(timeout=5.0)
|
|
finally:
|
|
self._restore_ws_data_dir()
|
|
|
|
def test_remove_account_leaves_tombstone_when_join_times_out(self):
|
|
"""F-4: remove_account must NOT pop the manager when the join times
|
|
out; it keeps a tombstone so a re-add reuses it instead of spawning a
|
|
duplicate worker."""
|
|
import webui_server as ws_module
|
|
_, ContinuousScrapeOrchestrator, _ = (
|
|
ws_module.PerAccountContinuousScrapeManager,
|
|
ws_module.ContinuousScrapeOrchestrator,
|
|
ws_module,
|
|
)
|
|
self._setup_ws_data_dir()
|
|
try:
|
|
orch = ContinuousScrapeOrchestrator()
|
|
aid = make_account_id()
|
|
create_account(TEST_DATA, aid, continuous_scraping={
|
|
"enabled": False, "interval_minutes": 1, "channels": [], "run_all_tracked": True,
|
|
})
|
|
mgr = orch._get_or_create(aid)
|
|
# Override join to return False (simulate a still-running thread).
|
|
original_join = mgr.join
|
|
mgr.join = lambda timeout=None: False
|
|
orch.remove_account(aid)
|
|
assert aid in orch.managers, "manager should be kept as tombstone"
|
|
assert orch.managers[aid]._removing is True
|
|
mgr.join = original_join
|
|
# Re-add the account: _get_or_create clears the tombstone and reuses
|
|
# the same manager object (no duplicate worker created).
|
|
mgr2 = orch._get_or_create(aid)
|
|
assert mgr2 is mgr, "re-add must reuse the tombstone manager"
|
|
assert mgr2._removing is False
|
|
finally:
|
|
self._restore_ws_data_dir()
|
|
|
|
|
|
# ── Cleanup all temp data ──────────────────────────────────────────────────
|
|
|
|
|
|
def cleanup_test_data():
|
|
if TEST_TMP.exists():
|
|
shutil.rmtree(str(TEST_TMP), ignore_errors=True)
|
|
|
|
|
|
import atexit # noqa: E402
|
|
atexit.register(cleanup_test_data) # noqa: E402
|