wehub-resource-sync
09e9f3545f
Test / Code Quality (push) Has been cancelled
Test / Test (macos-latest, Python 3.10) (push) Has been cancelled
Test / Test (macos-latest, Python 3.11) (push) Has been cancelled
Test / Test (macos-latest, Python 3.12) (push) Has been cancelled
Test / Test (macos-latest, Python 3.13) (push) Has been cancelled
Test / Test (macos-latest, Python 3.14) (push) Has been cancelled
Test / Test (ubuntu-latest, Python 3.10) (push) Has been cancelled
Test / Test (ubuntu-latest, Python 3.11) (push) Has been cancelled
Test / Test (ubuntu-latest, Python 3.12) (push) Has been cancelled
Test / Test (ubuntu-latest, Python 3.13) (push) Has been cancelled
Test / Test (ubuntu-latest, Python 3.14) (push) Has been cancelled
Test / Test (windows-latest, Python 3.10) (push) Has been cancelled
Test / Test (windows-latest, Python 3.11) (push) Has been cancelled
Test / Test (windows-latest, Python 3.12) (push) Has been cancelled
Test / Test (windows-latest, Python 3.13) (push) Has been cancelled
Test / Test (windows-latest, Python 3.14) (push) Has been cancelled
CodeQL / Analyze (push) Has been cancelled
dependency-audit / pip-audit (push) Has been cancelled
379 lines
16 KiB
Python
379 lines
16 KiB
Python
"""Tests for the per-loop / per-resolved-storage-path refresh lock registry.
|
||
|
||
The module-global ``_REFRESH_LOCK = asyncio.Lock()`` was removed because
|
||
``asyncio.Lock`` binds to the first event loop that uses it, breaking on
|
||
cross-loop / cross-thread usage. The replacement mirrors the keepalive
|
||
``_get_poke_lock`` pattern: a ``WeakKeyDictionary`` keyed on the running
|
||
loop, with per-resolved-storage-path inner locks. Cross-loop atomicity
|
||
for ``_REFRESH_GENERATIONS`` is provided by the sync ``_REFRESH_STATE_LOCK``.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import gc
|
||
import threading
|
||
from pathlib import Path
|
||
|
||
import pytest
|
||
|
||
from notebooklm import auth as auth_mod
|
||
from notebooklm._auth import refresh as _auth_refresh
|
||
|
||
|
||
@pytest.fixture(autouse=True)
|
||
def _clear_refresh_state():
|
||
"""Reset module state between tests."""
|
||
auth_mod._REFRESH_GENERATIONS.clear()
|
||
# The WeakKeyDictionary is mutated indirectly; we don't reset it here so
|
||
# the cleanup test can observe natural GC behavior.
|
||
yield
|
||
auth_mod._REFRESH_GENERATIONS.clear()
|
||
|
||
|
||
class TestPerLoopLockIdentity:
|
||
def test_different_loops_get_different_locks(self):
|
||
"""A lock created in loop X must NOT be returned to loop Y.
|
||
|
||
Both loops must be kept alive simultaneously so the WeakKeyDictionary
|
||
entries survive long enough to compare the locks by identity. Using
|
||
``asyncio.run`` for each loop in turn allows Python 3.13's GC to
|
||
reclaim the first loop's entry before the second is created, after
|
||
which ``id()`` of the second lock can collide with the recycled
|
||
memory address of the first.
|
||
"""
|
||
path = Path("/tmp/notebooklm-test/storage_state.json")
|
||
|
||
loop_a = asyncio.new_event_loop()
|
||
loop_b = asyncio.new_event_loop()
|
||
try:
|
||
|
||
async def _grab():
|
||
return auth_mod._get_refresh_lock(path)
|
||
|
||
lock_a = loop_a.run_until_complete(_grab())
|
||
lock_b = loop_b.run_until_complete(_grab())
|
||
# Hold strong references to both locks AND both loops until
|
||
# after the assertion so neither can be GC'd / address-recycled.
|
||
assert lock_a is not lock_b, (
|
||
"Different event loops must produce distinct asyncio.Lock instances"
|
||
)
|
||
finally:
|
||
loop_a.close()
|
||
loop_b.close()
|
||
|
||
def test_same_loop_same_path_returns_same_lock(self):
|
||
"""Repeated calls within one loop for the same path return the same lock."""
|
||
path = Path("/tmp/notebooklm-test/storage_state.json")
|
||
|
||
async def capture_two_locks():
|
||
return id(auth_mod._get_refresh_lock(path)), id(auth_mod._get_refresh_lock(path))
|
||
|
||
a, b = asyncio.run(capture_two_locks())
|
||
assert a == b
|
||
|
||
def test_same_loop_different_paths_get_different_locks(self):
|
||
"""Distinct storage paths within one loop get distinct locks."""
|
||
path_a = Path("/tmp/notebooklm-test/profile_a/storage_state.json")
|
||
path_b = Path("/tmp/notebooklm-test/profile_b/storage_state.json")
|
||
|
||
async def capture():
|
||
return (
|
||
id(auth_mod._get_refresh_lock(path_a)),
|
||
id(auth_mod._get_refresh_lock(path_b)),
|
||
)
|
||
|
||
a, b = asyncio.run(capture())
|
||
assert a != b
|
||
|
||
|
||
class TestResolvedPathEquivalence:
|
||
def test_none_with_profile_and_explicit_path_share_lock(self, monkeypatch, tmp_path):
|
||
"""``(None, profile="foo")`` and the explicit path resolve to the same key.
|
||
|
||
``_fetch_tokens_with_refresh`` computes
|
||
``refresh_storage_path = storage_path or get_storage_path(profile=profile)``
|
||
and keys the lock on that resolved Path, so two callers using either
|
||
form share the same lock.
|
||
"""
|
||
resolved = tmp_path / "profile_foo" / "storage_state.json"
|
||
|
||
async def capture():
|
||
# Both calls pass the SAME resolved Path object the production code
|
||
# would compute. The registry keys on Path equality, not identity.
|
||
lock_via_none = auth_mod._get_refresh_lock(resolved)
|
||
# Build an equivalent Path object (different instance, same value).
|
||
equivalent = Path(str(resolved))
|
||
lock_via_explicit = auth_mod._get_refresh_lock(equivalent)
|
||
return id(lock_via_none), id(lock_via_explicit)
|
||
|
||
a, b = asyncio.run(capture())
|
||
assert a == b, "Equal Path values must hash to the same registry slot"
|
||
|
||
def test_symlink_and_real_path_share_lock_via_refresh(self, monkeypatch, tmp_path):
|
||
"""Two different surface representations of the same physical file
|
||
(symlink vs canonical absolute path, and relative vs absolute) must
|
||
flow through ``_fetch_tokens_with_refresh``'s path canonicalization
|
||
and end up sharing a single lock / generation-key.
|
||
|
||
The production fix is ``.expanduser().resolve()`` applied to
|
||
``refresh_storage_path`` before it is used as a key. Without that
|
||
normalization, two callers referring to the same on-disk file by
|
||
different paths would each get their own lock — defeating the
|
||
cross-loop / cross-thread refresh coalescing.
|
||
"""
|
||
# Set up a real file plus a symlink pointing at it.
|
||
real_dir = tmp_path / "real"
|
||
real_dir.mkdir()
|
||
real_path = real_dir / "storage_state.json"
|
||
real_path.write_text('{"cookies": [], "origins": []}')
|
||
|
||
link_dir = tmp_path / "via_link"
|
||
link_dir.symlink_to(real_dir, target_is_directory=True)
|
||
symlinked_path = link_dir / "storage_state.json"
|
||
|
||
# Sanity: surface paths differ, resolved targets match.
|
||
assert symlinked_path != real_path
|
||
assert symlinked_path.resolve() == real_path.resolve()
|
||
|
||
captured_keys: list[str] = []
|
||
original_get_refresh_lock = auth_mod._get_refresh_lock
|
||
|
||
def spy_get_refresh_lock(p):
|
||
captured_keys.append(str(p))
|
||
return original_get_refresh_lock(p)
|
||
|
||
# Force a refresh attempt and short-circuit downstream side effects.
|
||
monkeypatch.setenv(auth_mod.NOTEBOOKLM_REFRESH_CMD_ENV, "dummy")
|
||
monkeypatch.setattr(_auth_refresh, "_get_refresh_lock", spy_get_refresh_lock)
|
||
|
||
call_phase = {"first": True}
|
||
|
||
async def fake_fetch_tokens_with_jar(jar, path, **kwargs):
|
||
if call_phase["first"]:
|
||
call_phase["first"] = False
|
||
raise ValueError("Authentication expired. Run 'notebooklm login'.")
|
||
return "csrf-token", "session-id"
|
||
|
||
async def fake_run_refresh_cmd(storage_path, profile):
|
||
return None
|
||
|
||
import httpx
|
||
|
||
def fake_build(_p):
|
||
return httpx.Cookies()
|
||
|
||
def fake_snapshot(_j):
|
||
return None
|
||
|
||
monkeypatch.setattr(_auth_refresh, "_fetch_tokens_with_jar", fake_fetch_tokens_with_jar)
|
||
monkeypatch.setattr(_auth_refresh, "_run_refresh_cmd", fake_run_refresh_cmd)
|
||
monkeypatch.setattr(_auth_refresh, "build_httpx_cookies_from_storage", fake_build)
|
||
monkeypatch.setattr(_auth_refresh, "snapshot_cookie_jar", fake_snapshot)
|
||
|
||
async def drive(path: Path):
|
||
jar = httpx.Cookies()
|
||
return await auth_mod._fetch_tokens_with_refresh(jar, storage_path=path)
|
||
|
||
# First caller uses the symlinked path.
|
||
asyncio.run(drive(symlinked_path))
|
||
# Reset phase so the second caller also triggers the refresh branch.
|
||
call_phase["first"] = True
|
||
# Second caller uses the canonical real path.
|
||
asyncio.run(drive(real_path))
|
||
|
||
# Both calls must have canonicalized to the same key.
|
||
assert len(captured_keys) == 2
|
||
assert captured_keys[0] == captured_keys[1], (
|
||
f"Symlinked and direct paths produced distinct keys: {captured_keys!r}"
|
||
)
|
||
# And the canonical key must equal the resolved real path.
|
||
assert captured_keys[0] == str(real_path.resolve())
|
||
|
||
|
||
class TestWeakKeyDictionaryCleanup:
|
||
def test_loops_are_garbage_collected(self):
|
||
"""When loops go out of scope, their inner dict is reclaimed."""
|
||
path = Path("/tmp/notebooklm-test/storage_state.json")
|
||
|
||
def spawn_and_drop():
|
||
"""Run a short-lived loop, return; do not retain a reference."""
|
||
|
||
async def inner():
|
||
auth_mod._get_refresh_lock(path)
|
||
|
||
asyncio.run(inner())
|
||
|
||
baseline_len = len(auth_mod._REFRESH_LOCKS_BY_LOOP)
|
||
for _ in range(5):
|
||
spawn_and_drop()
|
||
|
||
# Force collection of the closed loops.
|
||
gc.collect()
|
||
|
||
post_len = len(auth_mod._REFRESH_LOCKS_BY_LOOP)
|
||
# The WeakKeyDictionary should reclaim entries for collected loops.
|
||
# We assert the registry didn't grow unboundedly across iterations.
|
||
assert post_len <= baseline_len + 1, (
|
||
f"WeakKeyDictionary did not reclaim closed-loop entries "
|
||
f"(baseline={baseline_len}, post={post_len})"
|
||
)
|
||
|
||
|
||
class TestCrossLoopGenerationGuard:
|
||
def test_two_loops_at_most_two_refreshes(self, monkeypatch, tmp_path):
|
||
"""Two concurrent event loops both call ``_fetch_tokens_with_refresh``;
|
||
AT MOST 2 subprocess invocations occur (each loop runs once at the
|
||
worst), and crucially BOTH calls succeed without raising.
|
||
|
||
Race model: both loops observe an auth-expiry failure, each
|
||
capture the same pre-refresh generation, each acquire their OWN
|
||
per-loop asyncio lock (the registry hands out distinct locks per
|
||
loop), then race the check-and-claim under ``_REFRESH_STATE_LOCK``.
|
||
|
||
Legacy contract (before the gated-generation fix): this test
|
||
asserted ``run_count == 1`` because the old code bumped
|
||
``_REFRESH_GENERATIONS`` EAGERLY pre-subprocess; the cross-loop
|
||
loser saw the bump and skipped. That eager-bump behavior was the
|
||
root cause of the phantom-bump failure — when the subprocess
|
||
failed, the bump fooled concurrent waiters into skipping with
|
||
stale storage.
|
||
|
||
Current contract (gated-generation fix): generation is bumped
|
||
ONLY after the subprocess succeeds. Cross-loop callers cannot
|
||
signal "in flight" to each other (``asyncio.Future`` is
|
||
loop-bound). In the rare
|
||
cross-loop-concurrent-refresh case both loops may run their own
|
||
subprocess — equivalent to two ``RotateCookies`` POSTs against
|
||
the same storage. The end-state is correct (fresh cookies on
|
||
disk; last writer wins, but both write the same fresh data).
|
||
|
||
For SAME-LOOP coalescing (the dominant real-world case), the
|
||
per-loop in-flight future ensures exactly-once subprocess
|
||
execution; see ``tests/integration/concurrency/test_refresh_cmd_race.py``.
|
||
|
||
To force a deterministic race in a unit test, we use a barrier at
|
||
the point AFTER both threads have captured the generation but
|
||
BEFORE either has entered the inner sync mutex.
|
||
"""
|
||
import httpx
|
||
|
||
storage = tmp_path / "storage_state.json"
|
||
storage.write_text('{"cookies": [], "origins": []}')
|
||
|
||
# Force ``_should_try_refresh`` to return True.
|
||
monkeypatch.setenv(auth_mod.NOTEBOOKLM_REFRESH_CMD_ENV, "dummy")
|
||
|
||
run_count = 0
|
||
run_count_lock = threading.Lock()
|
||
|
||
# Barrier #1: align both fakes at the failure point so neither
|
||
# captures the generation before the other has had a chance to
|
||
# observe gen 0.
|
||
fail_barrier = threading.Barrier(2, timeout=5)
|
||
# Barrier #2: align both threads INSIDE their per-loop asyncio lock
|
||
# but BEFORE the inner sync-mutex check-and-claim. This guarantees
|
||
# both have already captured gen=0 pre-lock.
|
||
post_lock_barrier = threading.Barrier(2, timeout=5)
|
||
|
||
async def fake_run_refresh_cmd(storage_path, profile):
|
||
nonlocal run_count
|
||
with run_count_lock:
|
||
run_count += 1
|
||
await asyncio.sleep(0.01)
|
||
|
||
async def fake_fetch_tokens_with_jar(cookie_jar, storage_path, **kwargs):
|
||
if not getattr(cookie_jar, "_refresh_done", False):
|
||
cookie_jar._refresh_done = True
|
||
try:
|
||
fail_barrier.wait()
|
||
except threading.BrokenBarrierError:
|
||
pass
|
||
raise ValueError("Authentication expired. Run 'notebooklm login'.")
|
||
return "csrf-token", "session-id"
|
||
|
||
def fake_build_httpx_cookies(path):
|
||
return httpx.Cookies()
|
||
|
||
def fake_snapshot(jar):
|
||
return None
|
||
|
||
# Wrap ``_get_refresh_lock`` so we can interpose a barrier between
|
||
# the (pre-lock) generation capture and the (post-lock) inner
|
||
# sync-mutex check-and-claim.
|
||
original_get_refresh_lock = auth_mod._get_refresh_lock
|
||
|
||
class _BarrierLock:
|
||
def __init__(self, inner):
|
||
self._inner = inner
|
||
|
||
async def __aenter__(self):
|
||
await self._inner.acquire()
|
||
# Yield to a thread so the OTHER loop can also acquire its
|
||
# (distinct) per-loop lock before either touches the
|
||
# generation dict.
|
||
await asyncio.to_thread(post_lock_barrier.wait)
|
||
return self
|
||
|
||
async def __aexit__(self, exc_type, exc, tb):
|
||
self._inner.release()
|
||
return False
|
||
|
||
def wrapped_get_refresh_lock(p):
|
||
return _BarrierLock(original_get_refresh_lock(p))
|
||
|
||
monkeypatch.setattr(_auth_refresh, "_run_refresh_cmd", fake_run_refresh_cmd)
|
||
monkeypatch.setattr(_auth_refresh, "_fetch_tokens_with_jar", fake_fetch_tokens_with_jar)
|
||
monkeypatch.setattr(
|
||
_auth_refresh, "build_httpx_cookies_from_storage", fake_build_httpx_cookies
|
||
)
|
||
monkeypatch.setattr(_auth_refresh, "snapshot_cookie_jar", fake_snapshot)
|
||
monkeypatch.setattr(_auth_refresh, "_get_refresh_lock", wrapped_get_refresh_lock)
|
||
|
||
results: list[BaseException | tuple] = []
|
||
results_lock = threading.Lock()
|
||
|
||
def run_in_own_loop():
|
||
async def _work():
|
||
jar = httpx.Cookies()
|
||
return await auth_mod._fetch_tokens_with_refresh(jar, storage_path=storage)
|
||
|
||
try:
|
||
res = asyncio.run(_work())
|
||
with results_lock:
|
||
results.append(res)
|
||
except BaseException as exc: # noqa: BLE001
|
||
with results_lock:
|
||
results.append(exc)
|
||
|
||
thread_a = threading.Thread(target=run_in_own_loop)
|
||
thread_b = threading.Thread(target=run_in_own_loop)
|
||
thread_a.start()
|
||
thread_b.start()
|
||
thread_a.join(timeout=15)
|
||
thread_b.join(timeout=15)
|
||
|
||
assert not thread_a.is_alive() and not thread_b.is_alive(), (
|
||
"Threads failed to terminate; likely a deadlock"
|
||
)
|
||
# Both calls must have completed successfully.
|
||
assert len(results) == 2
|
||
for r in results:
|
||
assert not isinstance(r, BaseException), f"Refresh raised: {r!r}"
|
||
|
||
# The critical assertion under the gated-generation contract:
|
||
# AT MOST one subprocess invocation per loop (== 2 across two
|
||
# loops). Under the legacy eager-bump contract this asserted
|
||
# ``run_count == 1`` (cross-loop coalescing); the fix dropped
|
||
# eager-bump to close the phantom-bump failure mode, accepting
|
||
# that two cross-loop callers may both run their subprocess in
|
||
# the rare concurrent-refresh case (correct end-state: fresh
|
||
# cookies on disk).
|
||
assert 1 <= run_count <= 2, (
|
||
f"Expected 1–2 refresh invocations across loops, observed "
|
||
f"{run_count}. ``run_count == 0`` would mean the cross-loop "
|
||
"generation guard SKIPPED both refreshes (phantom-bump regression). "
|
||
"``run_count > 2`` means each loop ran refresh more than once "
|
||
"(per-loop coalescing broken)."
|
||
)
|