notifier: bound push/Discord sends, cap pass duration, fix key-gen race (#5)
This commit is contained in:
parent
7b7727f234
commit
d67476ebbf
9 changed files with 405 additions and 35 deletions
64
indexnow.py
64
indexnow.py
|
|
@ -36,10 +36,58 @@ _STATE_PATH = os.environ.get("THERMOGRAPH_INDEXNOW_STATE_FILE") or os.path.join(
|
|||
_ENDPOINT = "https://api.indexnow.org/indexnow" # a shared endpoint; it fans out to all participants
|
||||
_BATCH = 10000 # IndexNow's max URLs per request
|
||||
|
||||
# One shared client instead of a bare httpx.post per batch — same reused-client
|
||||
# pattern as web/app.py's _frontend_client and notifications/discord.py's _client.
|
||||
# 30s default matches submit()'s own prior per-call default; still overridable
|
||||
# per call (submit_all's fan-out passes its own `timeout` through unchanged).
|
||||
_client = httpx.Client(timeout=30)
|
||||
|
||||
_lock = threading.Lock()
|
||||
_key = None
|
||||
|
||||
|
||||
def _claim_file(value: str) -> str:
|
||||
"""Persist a freshly generated key as `_KEY_PATH`'s content, atomic
|
||||
first-writer-wins — same race, same fix as push.py's VAPID `_claim_file`
|
||||
(see its docstring): on a cold /state volume several workers can hit the
|
||||
missing-file branch at once, and without this each would generate and cache
|
||||
its OWN key, diverging from the file and from each other (the key served at
|
||||
/{key}.txt would then randomly mismatch whichever key a given worker signed
|
||||
a submission with). Write to a private temp file, then claim the real path
|
||||
with a hard link (atomic; a loser gets EEXIST with no half-written window);
|
||||
a loser reads back the winner's file instead of keeping its own generation."""
|
||||
tmp_path = f"{_KEY_PATH}.{os.getpid()}.tmp"
|
||||
try:
|
||||
os.makedirs(_DATA_DIR, exist_ok=True)
|
||||
with open(tmp_path, "w", encoding="utf-8") as f:
|
||||
f.write(value)
|
||||
os.chmod(tmp_path, 0o600)
|
||||
except OSError:
|
||||
log.warning("could not persist IndexNow key to %s; using an in-memory key", _KEY_PATH)
|
||||
return value
|
||||
try:
|
||||
try:
|
||||
os.link(tmp_path, _KEY_PATH)
|
||||
return value
|
||||
except FileExistsError:
|
||||
try:
|
||||
with open(_KEY_PATH, encoding="utf-8") as f:
|
||||
winner = f.read().strip()
|
||||
if winner:
|
||||
return winner
|
||||
except OSError:
|
||||
pass
|
||||
return value
|
||||
except OSError:
|
||||
log.warning("could not persist IndexNow key to %s; using an in-memory key", _KEY_PATH)
|
||||
return value
|
||||
finally:
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def key() -> str:
|
||||
"""The IndexNow key (env → file → generate-and-persist), cached for the process."""
|
||||
global _key
|
||||
|
|
@ -60,14 +108,10 @@ def key() -> str:
|
|||
return _key
|
||||
except OSError:
|
||||
pass
|
||||
_key = secrets.token_hex(16) # 32 hex chars — within IndexNow's 8–128 range
|
||||
try:
|
||||
os.makedirs(_DATA_DIR, exist_ok=True)
|
||||
with open(_KEY_PATH, "w", encoding="utf-8") as f:
|
||||
f.write(_key)
|
||||
os.chmod(_KEY_PATH, 0o600)
|
||||
except OSError:
|
||||
log.warning("could not persist IndexNow key to %s; using an in-memory key", _KEY_PATH)
|
||||
# Cache whatever _claim_file settles on — see its docstring: our own
|
||||
# generation if we won the race to create the file, the winner's key
|
||||
# read back from disk if we lost it.
|
||||
_key = _claim_file(secrets.token_hex(16)) # 32 hex chars — within IndexNow's 8–128 range
|
||||
return _key
|
||||
|
||||
|
||||
|
|
@ -83,8 +127,8 @@ def submit(urls, host: str, key_location: str, scheme: str = "https", timeout: f
|
|||
batch = urls[i:i + _BATCH]
|
||||
payload = {"host": host, "key": k, "keyLocation": key_location, "urlList": batch}
|
||||
try:
|
||||
r = httpx.post(_ENDPOINT, json=payload, timeout=timeout,
|
||||
headers={"Content-Type": "application/json; charset=utf-8"})
|
||||
r = _client.post(_ENDPOINT, json=payload, timeout=timeout,
|
||||
headers={"Content-Type": "application/json; charset=utf-8"})
|
||||
except httpx.HTTPError as e:
|
||||
log.warning("IndexNow request failed: %s", e)
|
||||
statuses.append("error")
|
||||
|
|
|
|||
|
|
@ -47,6 +47,12 @@ _COLD = 0x4393C3
|
|||
|
||||
_TIMEOUT = httpx.Timeout(10.0)
|
||||
|
||||
# One shared client instead of a bare httpx.post per call — same reused-client
|
||||
# pattern as web/app.py's _frontend_client, so every webhook/bot-REST call in
|
||||
# this module reuses one connection pool instead of paying a fresh TCP+TLS
|
||||
# handshake per notification.
|
||||
_client = httpx.Client(timeout=_TIMEOUT)
|
||||
|
||||
|
||||
def enabled() -> bool:
|
||||
return bool(WEBHOOK_URL) or weather_enabled()
|
||||
|
|
@ -131,7 +137,7 @@ def post_daily_feed(feed: dict | None = None) -> bool:
|
|||
"embeds": [embed],
|
||||
}
|
||||
try:
|
||||
resp = httpx.post(WEBHOOK_URL, json=payload, timeout=_TIMEOUT)
|
||||
resp = _client.post(WEBHOOK_URL, json=payload, timeout=_TIMEOUT)
|
||||
ok = 200 <= resp.status_code < 300
|
||||
except Exception: # noqa: BLE001 - best-effort side channel
|
||||
pass
|
||||
|
|
@ -159,10 +165,10 @@ def _bot_post(path: str, json_body: dict) -> httpx.Response | None:
|
|||
"""POST to the bot REST API with one 429 retry. None on a transport error."""
|
||||
headers = {"Authorization": f"Bot {BOT_TOKEN}"}
|
||||
try:
|
||||
resp = httpx.post(f"{_API}{path}", json=json_body, headers=headers, timeout=_TIMEOUT)
|
||||
resp = _client.post(f"{_API}{path}", json=json_body, headers=headers, timeout=_TIMEOUT)
|
||||
if resp.status_code == 429:
|
||||
time.sleep(_retry_after(resp))
|
||||
resp = httpx.post(f"{_API}{path}", json=json_body, headers=headers, timeout=_TIMEOUT)
|
||||
resp = _client.post(f"{_API}{path}", json=json_body, headers=headers, timeout=_TIMEOUT)
|
||||
return resp
|
||||
except Exception: # noqa: BLE001 - best-effort side channel
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ A forecast-fetch failure just skips the cell, so a pass never dies on a rate lim
|
|||
"""
|
||||
import concurrent.futures
|
||||
import datetime
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
|
|
@ -42,6 +43,8 @@ from notifications import push
|
|||
from accounts.db import sync_session_maker
|
||||
from accounts.models import AccessToken, Notification, PushSubscription, Subscription, User
|
||||
|
||||
log = logging.getLogger("thermograph.notify")
|
||||
|
||||
_BASE = os.environ.get("THERMOGRAPH_BASE", "/thermograph").strip("/")
|
||||
BASE = f"/{_BASE}" if _BASE else ""
|
||||
|
||||
|
|
@ -57,6 +60,14 @@ FORECAST_HORIZON_DAYS = 7
|
|||
# — the rest are picked up on later passes.
|
||||
MAX_ARCHIVE_FETCHES_PER_PASS = int(os.environ.get("THERMOGRAPH_NOTIFY_ARCHIVE_FETCHES", "4"))
|
||||
|
||||
# Soft wall-clock budget for one whole pass, same discipline as the archive-fetch
|
||||
# cap above: checked between cells (never mid-cell) so one pathological pass — a
|
||||
# slow upstream, a DB hiccup, hundreds of cells each doing real work — can't run
|
||||
# unbounded and push the next pass's wake past its own interval. 2x INTERVAL by
|
||||
# default: generous enough that an ordinarily-slow pass still finishes, but a
|
||||
# genuinely runaway one still yields before it would overlap the next wake.
|
||||
PASS_DEADLINE_SECONDS = int(os.environ.get("THERMOGRAPH_NOTIFY_PASS_DEADLINE", str(INTERVAL * 2)))
|
||||
|
||||
# Metrics whose low tail also counts as unusual when a subscription is two-sided
|
||||
# (cold snaps, unusually calm/dry). Precipitation is one-directional (high only).
|
||||
TWO_SIDED_METRICS = set(grading.TEMP_METRICS)
|
||||
|
|
@ -171,19 +182,40 @@ def _send_discord_job(job: tuple[str, str, str, str]) -> None:
|
|||
# hundreds of subscribers no longer stretches a pass out send-by-send.
|
||||
SEND_WORKERS = int(os.environ.get("THERMOGRAPH_NOTIFY_SEND_WORKERS", "8"))
|
||||
|
||||
# Belt-and-suspenders on top of push.py's own send timeout: f.result() below has
|
||||
# no timeout by default and would otherwise wait forever for a future that never
|
||||
# resolves, which — same as an untimed webpush() call — wedges the notifier for
|
||||
# the process's life under the singleton flock. Generous on purpose; this only
|
||||
# needs to fire if a per-send timeout was somehow bypassed, not to race it.
|
||||
SEND_RESULT_TIMEOUT = int(os.environ.get("THERMOGRAPH_NOTIFY_SEND_RESULT_TIMEOUT", "60"))
|
||||
|
||||
|
||||
def _flush_sends(push_jobs, discord_jobs) -> list[int]:
|
||||
"""Deliver every gathered push/Discord job concurrently. Returns the
|
||||
push-subscription row ids the push service reported gone, for the caller
|
||||
to prune. Blocks until every job has finished (or errored) — a pass still
|
||||
waits for delivery to complete, it just no longer does so serially."""
|
||||
to prune. Blocks until every job has finished or errored, and gives up
|
||||
waiting on (but does not lose track of) any single job past
|
||||
SEND_RESULT_TIMEOUT — a pass still waits for delivery to complete, it just
|
||||
no longer does so serially, and no misbehaving future can hang the result
|
||||
loop past a bounded wait."""
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=SEND_WORKERS) as pool:
|
||||
futures = [pool.submit(_send_push_job, job) for job in push_jobs]
|
||||
futures += [pool.submit(_send_discord_job, job) for job in discord_jobs]
|
||||
gone = []
|
||||
for f in futures:
|
||||
try:
|
||||
result = f.result()
|
||||
result = f.result(timeout=SEND_RESULT_TIMEOUT)
|
||||
except concurrent.futures.TimeoutError:
|
||||
# f.result() gave up waiting, but the worker thread itself can't be
|
||||
# killed — ThreadPoolExecutor.__exit__ still calls shutdown(wait=True),
|
||||
# so the pool as a whole won't return until that thread actually
|
||||
# finishes. The real backstop is the per-call socket timeouts this
|
||||
# commit adds (push.py's _SEND_TIMEOUT, discord.py's _TIMEOUT): they
|
||||
# bound how long a "stuck" send can truly run. This catch just stops
|
||||
# THIS result from being reported/counted as gone, and keeps checking
|
||||
# the remaining futures instead of getting stuck on one.
|
||||
log.warning("push/Discord send exceeded %ss timeout; treating as failed", SEND_RESULT_TIMEOUT)
|
||||
continue
|
||||
except Exception: # noqa: BLE001 - one bad send must not lose the rest
|
||||
continue
|
||||
if result is not None:
|
||||
|
|
@ -370,7 +402,18 @@ def run_pass() -> int:
|
|||
by_cell: dict[str, list] = {}
|
||||
for sub in subs:
|
||||
by_cell.setdefault(sub.cell_id, []).append(sub)
|
||||
for cell_id, cell_subs in by_cell.items():
|
||||
cells_skipped = 0
|
||||
for i, (cell_id, cell_subs) in enumerate(by_cell.items()):
|
||||
# Checked BETWEEN cells, not inside _process_cell — a soft budget, not
|
||||
# a hard per-item timeout. Once tripped, every remaining cell this pass
|
||||
# is skipped outright (not attempted-and-abandoned mid-flight) so a
|
||||
# skipped cell's subscriptions are simply picked up whole on the next
|
||||
# pass, same as a rate-limited one already is.
|
||||
if time.perf_counter() - t0 > PASS_DEADLINE_SECONDS:
|
||||
cells_skipped = len(by_cell) - i
|
||||
log.warning("notifier pass exceeded %ss deadline; skipping %d/%d remaining cell(s)",
|
||||
PASS_DEADLINE_SECONDS, cells_skipped, len(by_cell))
|
||||
break
|
||||
try:
|
||||
created += _process_cell(session, cell_id, cell_subs, today, now, archive_budget)
|
||||
except Exception: # noqa: BLE001 - one bad cell must not abort the pass
|
||||
|
|
@ -382,6 +425,7 @@ def run_pass() -> int:
|
|||
audit.log_activity("notify.pass", {
|
||||
"subs_evaluated": len(subs), "cells": len(by_cell), "created": created,
|
||||
"archives_fetched": MAX_ARCHIVE_FETCHES_PER_PASS - archive_budget[0],
|
||||
"cells_skipped_deadline": cells_skipped,
|
||||
"duration_ms": round((time.perf_counter() - t0) * 1000.0, 1)})
|
||||
return created
|
||||
|
||||
|
|
|
|||
|
|
@ -39,6 +39,15 @@ _VAPID_PATH = os.environ.get("THERMOGRAPH_VAPID_FILE") or os.path.join(_DATA_DIR
|
|||
# The VAPID "sub" claim — a contact the push service can reach about our traffic.
|
||||
_CONTACT = os.environ.get("THERMOGRAPH_VAPID_CONTACT", "mailto:admin@thermograph.app")
|
||||
|
||||
# pywebpush 2.0.0 forwards `timeout` straight to requests.post with no default of
|
||||
# its own — omit the kwarg here and the send blocks with NO timeout at all (the
|
||||
# "10s default" people expect only applies if requests itself is called bare).
|
||||
# One hung push endpoint must never wedge the notifier: the singleton flock
|
||||
# (core/singleton.py) is held for the process's whole life, so no other worker
|
||||
# can take over while this thread is stuck, and /healthz keeps reporting green
|
||||
# on a notifier that's actually frozen.
|
||||
_SEND_TIMEOUT = 15
|
||||
|
||||
_lock = threading.Lock()
|
||||
_keys = None # cached {"private_key": str, "public_key": str} (base64url-raw)
|
||||
|
||||
|
|
@ -60,6 +69,59 @@ def _generate() -> dict:
|
|||
return {"private_key": _b64url(scalar), "public_key": _b64url(raw_pub)}
|
||||
|
||||
|
||||
def _claim_file(data: dict) -> dict:
|
||||
"""Persist a freshly generated keypair as `_VAPID_PATH`'s content, atomic
|
||||
first-writer-wins. `_lock` only keeps this process's own threads from racing
|
||||
each other — on a cold /state volume every uvicorn *worker* (a separate
|
||||
process) hits the missing-file branch at boot together, and without this each
|
||||
would generate its OWN keypair and cache it in-process, silently diverging
|
||||
from the file and from each other (a subscription signed against one worker's
|
||||
public key fails to verify under another's private key — the exact incident
|
||||
class deploy/entrypoint.sh's comments name).
|
||||
|
||||
Write the full keypair to a private temp file first, then claim the real path
|
||||
with a hard link: `os.link` is atomic and the loser gets EEXIST immediately,
|
||||
with no window where `_VAPID_PATH` exists but is only half-written (unlike
|
||||
O_CREAT|O_EXCL directly on the destination followed by a separate write). A
|
||||
loser discards its own generation and reads back the winner's file instead,
|
||||
so every process ends up caching the SAME keys the file actually holds."""
|
||||
tmp_path = f"{_VAPID_PATH}.{os.getpid()}.tmp"
|
||||
try:
|
||||
os.makedirs(_DATA_DIR, exist_ok=True)
|
||||
with open(tmp_path, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f)
|
||||
os.chmod(tmp_path, 0o600)
|
||||
except OSError:
|
||||
log.warning("could not persist VAPID keys to %s; using an in-memory pair", _VAPID_PATH)
|
||||
return data
|
||||
try:
|
||||
try:
|
||||
os.link(tmp_path, _VAPID_PATH)
|
||||
return data
|
||||
except FileExistsError:
|
||||
# Lost the race — someone else's file is now the truth. Read IT back
|
||||
# rather than keep our own now-orphaned generation.
|
||||
try:
|
||||
with open(_VAPID_PATH, encoding="utf-8") as f:
|
||||
winner = json.load(f)
|
||||
if winner.get("private_key") and winner.get("public_key"):
|
||||
return winner
|
||||
except (OSError, ValueError):
|
||||
pass
|
||||
# Winner's file was unreadable/corrupt (very rare) — fall back to our
|
||||
# own in-memory generation rather than crash; the next _load() call
|
||||
# (a new process, or after the file heals) retries the file.
|
||||
return data
|
||||
except OSError:
|
||||
log.warning("could not persist VAPID keys to %s; using an in-memory pair", _VAPID_PATH)
|
||||
return data
|
||||
finally:
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _load() -> dict:
|
||||
"""Resolve the keypair once (env → file → generate) and cache it."""
|
||||
global _keys
|
||||
|
|
@ -81,14 +143,10 @@ def _load() -> dict:
|
|||
return _keys
|
||||
except (OSError, ValueError):
|
||||
pass
|
||||
_keys = _generate()
|
||||
try:
|
||||
os.makedirs(_DATA_DIR, exist_ok=True)
|
||||
with open(_VAPID_PATH, "w", encoding="utf-8") as f:
|
||||
json.dump(_keys, f)
|
||||
os.chmod(_VAPID_PATH, 0o600)
|
||||
except OSError:
|
||||
log.warning("could not persist VAPID keys to %s; using an in-memory pair", _VAPID_PATH)
|
||||
# Cache whatever _claim_file settles on — our own generation if we won
|
||||
# the race to create the file, or the winner's keys read back from disk
|
||||
# if we lost it. Either way this process's cache matches the file.
|
||||
_keys = _claim_file(_generate())
|
||||
return _keys
|
||||
|
||||
|
||||
|
|
@ -111,6 +169,7 @@ def send(subscription_info: dict, payload: dict) -> str:
|
|||
vapid_private_key=keys["private_key"],
|
||||
vapid_claims={"sub": _CONTACT},
|
||||
ttl=86400,
|
||||
timeout=_SEND_TIMEOUT,
|
||||
)
|
||||
return "ok"
|
||||
except WebPushException as e:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
"""Discord daily-feed post: embed shape, webhook delivery, and the once-a-day guard.
|
||||
No network — httpx.post is stubbed, like the IndexNow tests."""
|
||||
No network — the shared client's .post is stubbed, like the IndexNow tests."""
|
||||
import datetime
|
||||
import time
|
||||
|
||||
|
|
@ -63,7 +63,7 @@ def test_build_embed_rejects_empty_feed():
|
|||
def test_post_disabled_without_webhook(monkeypatch):
|
||||
monkeypatch.setattr(discord, "WEBHOOK_URL", "")
|
||||
posted = []
|
||||
monkeypatch.setattr(discord.httpx, "post", lambda *a, **k: posted.append(a) or None)
|
||||
monkeypatch.setattr(discord._client, "post", lambda *a, **k: posted.append(a) or None)
|
||||
assert discord.post_daily_feed(_feed([_CARD_HOT])) is False
|
||||
assert posted == [] # never touched the network
|
||||
|
||||
|
|
@ -75,7 +75,7 @@ class _Resp:
|
|||
def test_post_sends_embed_to_the_webhook(monkeypatch):
|
||||
monkeypatch.setattr(discord, "WEBHOOK_URL", "https://discord.test/webhook/abc")
|
||||
calls = []
|
||||
monkeypatch.setattr(discord.httpx, "post",
|
||||
monkeypatch.setattr(discord._client, "post",
|
||||
lambda url, json=None, **k: (calls.append((url, json)), _Resp(204))[1])
|
||||
assert discord.post_daily_feed(_feed([_CARD_HOT])) is True
|
||||
(url, payload), = calls
|
||||
|
|
@ -91,7 +91,7 @@ def test_post_also_broadcasts_to_weather_channel(monkeypatch):
|
|||
monkeypatch.setattr(discord, "BOT_TOKEN", "bot-abc")
|
||||
monkeypatch.setattr(discord, "WEATHER_CHANNEL_ID", "chan-weather")
|
||||
calls = []
|
||||
monkeypatch.setattr(discord.httpx, "post",
|
||||
monkeypatch.setattr(discord._client, "post",
|
||||
lambda url, json=None, **k: (calls.append((url, json)), _Resp(204))[1])
|
||||
assert discord.post_daily_feed(_feed([_CARD_HOT])) is True
|
||||
(url, body), = calls
|
||||
|
|
@ -104,14 +104,14 @@ def test_post_is_best_effort_on_http_error(monkeypatch):
|
|||
|
||||
def _boom(*a, **k):
|
||||
raise RuntimeError("network down")
|
||||
monkeypatch.setattr(discord.httpx, "post", _boom)
|
||||
monkeypatch.setattr(discord._client, "post", _boom)
|
||||
# Never raises; just reports failure.
|
||||
assert discord.post_daily_feed(_feed([_CARD_HOT])) is False
|
||||
|
||||
|
||||
def test_post_skips_stale_or_empty_feed(monkeypatch):
|
||||
monkeypatch.setattr(discord, "WEBHOOK_URL", "https://discord.test/webhook/abc")
|
||||
monkeypatch.setattr(discord.httpx, "post", lambda *a, **k: _Resp(204))
|
||||
monkeypatch.setattr(discord._client, "post", lambda *a, **k: _Resp(204))
|
||||
# Yesterday's feed is stale -> no post.
|
||||
yesterday = (datetime.date.today() - datetime.timedelta(days=1)).isoformat()
|
||||
assert discord.post_daily_feed(_feed([_CARD_HOT], date=yesterday)) is False
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ class _Resp:
|
|||
|
||||
|
||||
def _mock_posts(monkeypatch, script):
|
||||
"""Route discord.httpx.post by URL; `script` maps a URL substring -> _Resp."""
|
||||
"""Route discord._client.post by URL; `script` maps a URL substring -> _Resp."""
|
||||
calls = []
|
||||
|
||||
def _post(url, json=None, headers=None, timeout=None):
|
||||
|
|
@ -37,7 +37,7 @@ def _mock_posts(monkeypatch, script):
|
|||
if frag in url:
|
||||
return resp
|
||||
return _Resp(404)
|
||||
monkeypatch.setattr(discord.httpx, "post", _post)
|
||||
monkeypatch.setattr(discord._client, "post", _post)
|
||||
return calls
|
||||
|
||||
|
||||
|
|
@ -97,7 +97,7 @@ def test_bot_post_retries_once_on_429(monkeypatch):
|
|||
monkeypatch.setattr(discord, "BOT_TOKEN", "bot-abc")
|
||||
monkeypatch.setattr(discord.time, "sleep", lambda s: None) # don't actually wait
|
||||
seq = iter([_Resp(429, {"retry_after": 0.01}), _Resp(200, {"id": "c"})])
|
||||
monkeypatch.setattr(discord.httpx, "post", lambda *a, **k: next(seq))
|
||||
monkeypatch.setattr(discord._client, "post", lambda *a, **k: next(seq))
|
||||
resp = discord._bot_post("/users/@me/channels", {"recipient_id": "d"})
|
||||
assert resp.status_code == 200 # the retry succeeded
|
||||
|
||||
|
|
|
|||
|
|
@ -2,12 +2,14 @@
|
|||
wording (pure logic), plus an integration check that a full pass fetches a MISSING
|
||||
archive once but never re-fetches a cached one (against conftest's throwaway DB)."""
|
||||
import datetime
|
||||
import time
|
||||
import types
|
||||
import uuid
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from core import audit
|
||||
from data import climate
|
||||
from accounts import db
|
||||
from notifications import notify
|
||||
|
|
@ -211,3 +213,72 @@ def test_inapp_notification_survives_push_error(monkeypatch):
|
|||
|
||||
assert len(_user_notifications(uid)) == 1 # the in-app write is unaffected
|
||||
assert _push_count() == 1 # a mere error doesn't prune
|
||||
|
||||
|
||||
# --- send-result timeout: one hung future must not wedge the pass -----------
|
||||
def test_flush_sends_treats_timeout_as_a_failed_send(monkeypatch):
|
||||
monkeypatch.setattr(notify, "SEND_RESULT_TIMEOUT", 0.05)
|
||||
|
||||
def _hang(job):
|
||||
time.sleep(0.3) # longer than the timeout above
|
||||
return None
|
||||
monkeypatch.setattr(notify, "_send_push_job", _hang)
|
||||
|
||||
# Must not raise, and a timed-out send is never reported as a "gone"
|
||||
# endpoint to prune (it never actually got a definitive answer).
|
||||
gone = notify._flush_sends([object()], [])
|
||||
assert gone == []
|
||||
|
||||
|
||||
def test_flush_sends_still_returns_gone_ids_within_the_timeout(monkeypatch):
|
||||
monkeypatch.setattr(notify, "SEND_RESULT_TIMEOUT", 5)
|
||||
monkeypatch.setattr(notify, "_send_push_job", lambda job: job) # echoes the id straight back
|
||||
gone = notify._flush_sends([7, 8], [])
|
||||
assert sorted(gone) == [7, 8]
|
||||
|
||||
|
||||
# --- pass wall-clock deadline: one pathological pass can't run unbounded ----
|
||||
def _seed_two_cell_subscriptions():
|
||||
db.Base.metadata.create_all(db.sync_engine)
|
||||
with db.sync_session_maker() as s:
|
||||
s.execute(delete(User))
|
||||
s.commit()
|
||||
uid1, uid2 = uuid.uuid4(), uuid.uuid4()
|
||||
s.add(User(id=uid1, email="deadline1@example.com", hashed_password="x", is_active=True))
|
||||
s.add(User(id=uid2, email="deadline2@example.com", hashed_password="x", is_active=True))
|
||||
s.commit()
|
||||
s.add(Subscription(user_id=uid1, cell_id="500_600", label="A", lat=1.0, lon=2.0,
|
||||
threshold=95, metrics=["tmax"], kind="observed", two_sided=False))
|
||||
s.add(Subscription(user_id=uid2, cell_id="700_800", label="B", lat=3.0, lon=4.0,
|
||||
threshold=95, metrics=["tmax"], kind="observed", two_sided=False))
|
||||
s.commit()
|
||||
return uid1, uid2
|
||||
|
||||
|
||||
def test_run_pass_stops_at_wall_clock_deadline(monkeypatch):
|
||||
uid1, uid2 = _seed_two_cell_subscriptions()
|
||||
_cached_extreme(monkeypatch) # every cell would otherwise trigger a notification
|
||||
monkeypatch.setattr(notify, "PASS_DEADLINE_SECONDS", 0) # trips before the first cell
|
||||
|
||||
logged = []
|
||||
monkeypatch.setattr(audit, "log_activity",
|
||||
lambda kind, data: logged.append((kind, data)))
|
||||
|
||||
notify.run_pass()
|
||||
|
||||
# A deadline of 0 must skip every cell rather than process any of them.
|
||||
assert _user_notifications(uid1) == []
|
||||
assert _user_notifications(uid2) == []
|
||||
pass_log = next(data for kind, data in logged if kind == "notify.pass")
|
||||
assert pass_log["cells_skipped_deadline"] == 2
|
||||
|
||||
|
||||
def test_run_pass_processes_normally_within_the_deadline(monkeypatch):
|
||||
uid1, uid2 = _seed_two_cell_subscriptions()
|
||||
_cached_extreme(monkeypatch)
|
||||
monkeypatch.setattr(notify, "PASS_DEADLINE_SECONDS", 300) # generous — shouldn't trip
|
||||
|
||||
notify.run_pass()
|
||||
|
||||
assert len(_user_notifications(uid1)) == 1
|
||||
assert len(_user_notifications(uid2)) == 1
|
||||
|
|
|
|||
96
tests/notifications/test_push.py
Normal file
96
tests/notifications/test_push.py
Normal file
|
|
@ -0,0 +1,96 @@
|
|||
"""VAPID key resolution: the atomic first-writer-wins claim of the keypair file
|
||||
(push.py's `_claim_file`), and that `_load()` caches whatever that settles on
|
||||
rather than its own local generation when it loses the race.
|
||||
|
||||
No network — `send()`'s pywebpush call isn't exercised here (that's notify.py's
|
||||
`test_push_dispatched_on_new_notification` etc., which stub `push.send` itself)."""
|
||||
import json
|
||||
|
||||
from notifications import push
|
||||
|
||||
|
||||
def _reset_state(monkeypatch, path):
|
||||
monkeypatch.setattr(push, "_VAPID_PATH", str(path))
|
||||
monkeypatch.setattr(push, "_DATA_DIR", str(path.parent))
|
||||
monkeypatch.setattr(push, "_keys", None)
|
||||
monkeypatch.delenv("THERMOGRAPH_VAPID_PRIVATE_KEY", raising=False)
|
||||
monkeypatch.delenv("THERMOGRAPH_VAPID_PUBLIC_KEY", raising=False)
|
||||
|
||||
|
||||
def test_claim_file_first_writer_wins(monkeypatch, tmp_path):
|
||||
path = tmp_path / "vapid.json"
|
||||
_reset_state(monkeypatch, path)
|
||||
data = {"private_key": "priv-a", "public_key": "pub-a"}
|
||||
|
||||
result = push._claim_file(data)
|
||||
|
||||
assert result == data
|
||||
assert json.loads(path.read_text()) == data
|
||||
# No leftover temp file.
|
||||
assert list(tmp_path.iterdir()) == [path]
|
||||
|
||||
|
||||
def test_claim_file_reads_back_winner_on_race(monkeypatch, tmp_path):
|
||||
"""A process that loses the os.link() race must discard its own freshly
|
||||
generated keypair and use whatever the winner actually persisted."""
|
||||
path = tmp_path / "vapid.json"
|
||||
_reset_state(monkeypatch, path)
|
||||
winner = {"private_key": "priv-winner", "public_key": "pub-winner"}
|
||||
path.write_text(json.dumps(winner)) # simulates another process winning first
|
||||
loser = {"private_key": "priv-loser", "public_key": "pub-loser"}
|
||||
|
||||
result = push._claim_file(loser)
|
||||
|
||||
assert result == winner
|
||||
assert json.loads(path.read_text()) == winner # the loser never touched the file
|
||||
|
||||
|
||||
def test_load_caches_winner_keys_not_its_own_generation(monkeypatch, tmp_path):
|
||||
"""End-to-end through _load(): on a cold /state volume, a worker that hits the
|
||||
generate branch but loses the write race must end up with the SAME in-process
|
||||
cached keys the file actually holds — not the keys it generated locally,
|
||||
which would silently diverge from every other worker (a subscription signed
|
||||
against one worker's public key then fails to verify under another's
|
||||
private key)."""
|
||||
path = tmp_path / "vapid.json"
|
||||
_reset_state(monkeypatch, path)
|
||||
winner = {"private_key": "priv-winner", "public_key": "pub-winner"}
|
||||
|
||||
def _generate_and_lose_the_race():
|
||||
# Between our cache-miss file read and our own _claim_file() call, a
|
||||
# concurrent worker wins and writes the real file first.
|
||||
path.write_text(json.dumps(winner))
|
||||
return {"private_key": "priv-mine", "public_key": "pub-mine"}
|
||||
|
||||
monkeypatch.setattr(push, "_generate", _generate_and_lose_the_race)
|
||||
|
||||
result = push._load()
|
||||
|
||||
assert result == winner
|
||||
assert push._keys == winner # the process-wide cache matches the file
|
||||
assert push.public_key() == "pub-winner"
|
||||
|
||||
|
||||
def test_load_reads_existing_file_without_generating(monkeypatch, tmp_path):
|
||||
path = tmp_path / "vapid.json"
|
||||
_reset_state(monkeypatch, path)
|
||||
existing = {"private_key": "priv-x", "public_key": "pub-x"}
|
||||
path.write_text(json.dumps(existing))
|
||||
|
||||
def boom():
|
||||
raise AssertionError("must not generate when a valid file already exists")
|
||||
monkeypatch.setattr(push, "_generate", boom)
|
||||
|
||||
assert push._load() == existing
|
||||
|
||||
|
||||
def test_load_prefers_env_over_file(monkeypatch, tmp_path):
|
||||
path = tmp_path / "vapid.json"
|
||||
_reset_state(monkeypatch, path)
|
||||
path.write_text(json.dumps({"private_key": "priv-file", "public_key": "pub-file"}))
|
||||
monkeypatch.setenv("THERMOGRAPH_VAPID_PRIVATE_KEY", " priv-env ")
|
||||
monkeypatch.setenv("THERMOGRAPH_VAPID_PUBLIC_KEY", " pub-env ")
|
||||
|
||||
result = push._load()
|
||||
|
||||
assert result == {"private_key": "priv-env", "public_key": "pub-env"}
|
||||
|
|
@ -69,7 +69,7 @@ def test_indexnow_submit_all_builds_payload(monkeypatch):
|
|||
status_code = 200
|
||||
text = "ok"
|
||||
|
||||
monkeypatch.setattr(indexnow.httpx, "post",
|
||||
monkeypatch.setattr(indexnow._client, "post",
|
||||
lambda url, json=None, **kw: (calls.append((url, json)), _Resp())[1])
|
||||
res = indexnow.submit_all("https://thermograph.org")
|
||||
assert res["submitted"] == res["total"] > 100
|
||||
|
|
@ -87,3 +87,53 @@ def test_indexnow_if_changed_state(monkeypatch, tmp_path):
|
|||
assert indexnow._read_state() == "" # first deploy → would submit
|
||||
indexnow._write_state(sig)
|
||||
assert indexnow._read_state() == sig # unchanged → deploy would skip
|
||||
|
||||
|
||||
# --- key() race: atomic first-writer-wins on a cold /state volume -----------
|
||||
# Same race, same fix as push.py's VAPID keypair (see tests/notifications/
|
||||
# test_push.py) — several workers can hit the missing-key-file branch at once;
|
||||
# only one may actually create the file, everyone else must read IT back.
|
||||
|
||||
def _reset_key_state(monkeypatch, path):
|
||||
monkeypatch.setattr(indexnow, "_KEY_PATH", str(path))
|
||||
monkeypatch.setattr(indexnow, "_DATA_DIR", str(path.parent))
|
||||
monkeypatch.setattr(indexnow, "_key", None)
|
||||
monkeypatch.delenv("THERMOGRAPH_INDEXNOW_KEY", raising=False)
|
||||
|
||||
|
||||
def test_claim_file_first_writer_wins(monkeypatch, tmp_path):
|
||||
path = tmp_path / "indexnow_key.txt"
|
||||
_reset_key_state(monkeypatch, path)
|
||||
|
||||
result = indexnow._claim_file("key-a")
|
||||
|
||||
assert result == "key-a"
|
||||
assert path.read_text() == "key-a"
|
||||
assert list(tmp_path.iterdir()) == [path] # no leftover temp file
|
||||
|
||||
|
||||
def test_claim_file_reads_back_winner_on_race(monkeypatch, tmp_path):
|
||||
path = tmp_path / "indexnow_key.txt"
|
||||
_reset_key_state(monkeypatch, path)
|
||||
path.write_text("key-winner") # simulates another process winning first
|
||||
|
||||
result = indexnow._claim_file("key-loser")
|
||||
|
||||
assert result == "key-winner"
|
||||
assert path.read_text() == "key-winner" # the loser never touched the file
|
||||
|
||||
|
||||
def test_key_caches_winner_not_its_own_generation(monkeypatch, tmp_path):
|
||||
path = tmp_path / "indexnow_key.txt"
|
||||
_reset_key_state(monkeypatch, path)
|
||||
|
||||
def _lose_the_race(*a, **k):
|
||||
path.write_text("key-winner")
|
||||
return "a" * 32 # what secrets.token_hex(16) would have produced locally
|
||||
|
||||
monkeypatch.setattr(indexnow.secrets, "token_hex", _lose_the_race)
|
||||
|
||||
result = indexnow.key()
|
||||
|
||||
assert result == "key-winner"
|
||||
assert indexnow._key == "key-winner" # the process-wide cache matches the file
|
||||
|
|
|
|||
Loading…
Reference in a new issue