Initial commit: grok-free-register-oss
Open-source Grok free registration CLI, xai_enroller auth pipeline, local auth service, tests and docs.
This commit is contained in:
chaos
0
tests/__init__.py
Normal file
0
tests/__init__.py
Normal file
92
tests/conftest.py
Normal file
92
tests/conftest.py
Normal file
@@ -0,0 +1,92 @@
|
||||
"""
|
||||
共享 fixtures — 供四层测试共用
|
||||
"""
|
||||
import sys
|
||||
import os
|
||||
import asyncio
|
||||
import pytest
|
||||
|
||||
# 确保项目根目录在 sys.path 中
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from grok_register.core.inventory import Inventory
|
||||
from grok_register.core.observer import Metrics
|
||||
from tests.fakes import (
|
||||
FakeTurnstile, FakeEmailService, FakeRegisterAPI,
|
||||
Conservation, EventLog,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def metrics():
|
||||
return Metrics()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def inventory(metrics):
|
||||
return Inventory(metrics=metrics)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sems():
|
||||
"""标准 Semaphore 集合(小容量,便于测试)。"""
|
||||
return {
|
||||
'physical': asyncio.Semaphore(4),
|
||||
't_slot': asyncio.Semaphore(4),
|
||||
'q_slot': asyncio.Semaphore(4),
|
||||
'q_pending': asyncio.Semaphore(6),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def large_sems():
|
||||
"""大容量 Semaphore(压力测试用)。"""
|
||||
return {
|
||||
'physical': asyncio.Semaphore(32),
|
||||
't_slot': asyncio.Semaphore(64),
|
||||
'q_slot': asyncio.Semaphore(64),
|
||||
'q_pending': asyncio.Semaphore(48),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def conservation(sems, inventory):
|
||||
return Conservation(
|
||||
sems['t_slot'], sems['q_slot'], sems['q_pending'], inventory
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_turnstile():
|
||||
return FakeTurnstile()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_email():
|
||||
return FakeEmailService()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_register():
|
||||
return FakeRegisterAPI()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def event_log():
|
||||
return EventLog()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stop():
|
||||
return asyncio.Event()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def file_lock():
|
||||
return asyncio.Lock()
|
||||
|
||||
|
||||
# ── 跳过需要网络的测试(本地开发时) ──
|
||||
def pytest_configure(config):
|
||||
config.addinivalue_line("markers", "slow: 压力测试(较慢)")
|
||||
config.addinivalue_line("markers", "needs_network: 需要外部网络")
|
||||
317
tests/fakes.py
Normal file
317
tests/fakes.py
Normal file
@@ -0,0 +1,317 @@
|
||||
"""
|
||||
Fake 服务 + 测试工具
|
||||
|
||||
FakeTurnstile — 可控 T 生成器:可注入延迟、失败、取消点
|
||||
FakeEmailService — 可控 Q 提供者:可注入延迟、超时、验证码
|
||||
FakeConsumer — 可控 C 消费器:可注入延迟、失败
|
||||
Conservation — slot 守恒断言器
|
||||
CancelPoint — 在指定 await 边界注入取消
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
import dataclasses
|
||||
from typing import Optional, Callable, Any
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# Fake T 生成器 (替代 solve_one_turnstile)
|
||||
# ═══════════════════════════════════════════
|
||||
class FakeTurnstile:
|
||||
"""可控 token 生成器。
|
||||
|
||||
config:
|
||||
delay: 生成延迟(秒)
|
||||
fail_rate: 0.0~1.0,随机失败概率
|
||||
token_seq: 预设 token 序列,None 则自动递增
|
||||
"""
|
||||
|
||||
def __init__(self, delay=0.01, fail_rate=0.0, token_seq=None):
|
||||
self.delay = delay
|
||||
self._fail_rate = fail_rate
|
||||
self._seq = token_seq or []
|
||||
self._idx = 0
|
||||
self._call_count = 0
|
||||
self._fail_count = 0
|
||||
self._cancel_at: Optional[int] = None # 第 N 次调用时注入取消
|
||||
|
||||
def set_cancel_at(self, n: int):
|
||||
"""第 n 次调用时注入 CancelledError。"""
|
||||
self._cancel_at = n
|
||||
|
||||
async def generate(self) -> Optional[str]:
|
||||
"""生成一个 token。返回 None 表示失败,raise CancelledError 表示取消。"""
|
||||
self._call_count += 1
|
||||
|
||||
if self._cancel_at is not None and self._call_count == self._cancel_at:
|
||||
raise asyncio.CancelledError(f'FakeTurnstile cancel at call {self._call_count}')
|
||||
|
||||
await asyncio.sleep(self.delay)
|
||||
|
||||
if self._fail_count < self._fail_rate * self._call_count:
|
||||
return None
|
||||
# 简单随机: fail_rate 作为概率
|
||||
import random
|
||||
if random.random() < self._fail_rate:
|
||||
self._fail_count += 1
|
||||
return None
|
||||
|
||||
if self._idx < len(self._seq):
|
||||
tok = self._seq[self._idx]
|
||||
self._idx += 1
|
||||
return tok
|
||||
self._idx += 1
|
||||
return f'fake_token_{self._idx}'
|
||||
|
||||
@property
|
||||
def call_count(self):
|
||||
return self._call_count
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# Fake 邮箱服务 (替代 create_email + poll_code)
|
||||
# ═══════════════════════════════════════════
|
||||
@dataclasses.dataclass
|
||||
class FakeEmailResult:
|
||||
handle: str
|
||||
email: str
|
||||
password: str
|
||||
code: Optional[str] = None # None = 轮询超时
|
||||
create_delay: float = 0.01
|
||||
poll_delay: float = 0.01
|
||||
|
||||
|
||||
class FakeEmailService:
|
||||
"""可控邮箱 + 验证码服务。
|
||||
|
||||
results: 预设结果队列,每次 create+poll 消费一个。
|
||||
"""
|
||||
|
||||
def __init__(self, results=None):
|
||||
self._results = list(results or [])
|
||||
self._idx = 0
|
||||
self._create_count = 0
|
||||
self._poll_count = 0
|
||||
self._cancel_create_at: Optional[int] = None
|
||||
self._cancel_poll_at: Optional[int] = None
|
||||
|
||||
def set_cancel_create_at(self, n: int):
|
||||
self._cancel_create_at = n
|
||||
|
||||
def set_cancel_poll_at(self, n: int):
|
||||
self._cancel_poll_at = n
|
||||
|
||||
async def create_email(self):
|
||||
"""创建邮箱。返回 (handle, email, password)。"""
|
||||
self._create_count += 1
|
||||
if self._cancel_create_at and self._create_count == self._cancel_create_at:
|
||||
raise asyncio.CancelledError(f'FakeEmail.create cancel at {self._create_count}')
|
||||
|
||||
if self._idx < len(self._results):
|
||||
r = self._results[self._idx]
|
||||
await asyncio.sleep(r.create_delay)
|
||||
return r.handle, r.email, r.password
|
||||
self._idx += 1
|
||||
n = self._idx
|
||||
await asyncio.sleep(0.01)
|
||||
return f'handle_{n}', f'user_{n}@test.com', f'pass_{n}'
|
||||
|
||||
async def poll_code(self, handle: str, max_wait=90):
|
||||
"""轮询验证码。返回 code 或 None。"""
|
||||
self._poll_count += 1
|
||||
if self._cancel_poll_at and self._poll_count == self._cancel_poll_at:
|
||||
raise asyncio.CancelledError(f'FakeEmail.poll cancel at {self._poll_count}')
|
||||
|
||||
# 找匹配的预设结果
|
||||
for r in self._results:
|
||||
if r.handle == handle:
|
||||
await asyncio.sleep(r.poll_delay)
|
||||
return r.code
|
||||
await asyncio.sleep(0.01)
|
||||
return '000000'
|
||||
|
||||
@property
|
||||
def create_count(self):
|
||||
return self._create_count
|
||||
|
||||
@property
|
||||
def poll_count(self):
|
||||
return self._poll_count
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# Fake 注册 API (替代 grpc_verify + server_action_register)
|
||||
# ═══════════════════════════════════════════
|
||||
class FakeRegisterAPI:
|
||||
"""可控注册 API。"""
|
||||
|
||||
def __init__(self, success_rate=1.0, delay=0.01):
|
||||
self.success_rate = success_rate
|
||||
self.delay = delay
|
||||
self._register_count = 0
|
||||
self._cancel_at: Optional[int] = None
|
||||
self._fail_at: Optional[int] = None
|
||||
|
||||
def set_cancel_at(self, n: int):
|
||||
self._cancel_at = n
|
||||
|
||||
def set_fail_at(self, n: int):
|
||||
self._fail_at = n
|
||||
|
||||
async def register(self, email, password, code, token) -> Optional[str]:
|
||||
"""执行注册。返回 SSO 或 None。"""
|
||||
self._register_count += 1
|
||||
if self._cancel_at and self._register_count == self._cancel_at:
|
||||
raise asyncio.CancelledError(f'FakeRegister cancel at {self._register_count}')
|
||||
if self._fail_at and self._register_count == self._fail_at:
|
||||
return None
|
||||
await asyncio.sleep(self.delay)
|
||||
import random
|
||||
if random.random() < self.success_rate:
|
||||
return f'sso_{email}_{int(time.time())}'
|
||||
return None
|
||||
|
||||
@property
|
||||
def register_count(self):
|
||||
return self._register_count
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# Slot 守恒断言器
|
||||
# ═══════════════════════════════════════════
|
||||
class Conservation:
|
||||
"""检查 T/Q slot 守恒不变量。
|
||||
|
||||
守恒公式:
|
||||
slot_total = free_slots + inventory_depth + pairleased_held + worker_held
|
||||
"""
|
||||
|
||||
def __init__(self, t_slot_sem, q_slot_sem, q_pending_sem, inventory):
|
||||
self.t_slot_sem = t_slot_sem
|
||||
self.q_slot_sem = q_slot_sem
|
||||
self.q_pending_sem = q_pending_sem
|
||||
self.inventory = inventory
|
||||
# 追踪 Worker 持有的未发布资源
|
||||
self._worker_held_t = 0
|
||||
self._worker_held_q = 0
|
||||
|
||||
def worker_acquire_t(self):
|
||||
self._worker_held_t += 1
|
||||
|
||||
def worker_release_t(self):
|
||||
self._worker_held_t -= 1
|
||||
|
||||
def worker_acquire_q(self):
|
||||
self._worker_held_q += 1
|
||||
|
||||
def worker_release_q(self):
|
||||
self._worker_held_q -= 1
|
||||
|
||||
def check_t_conservation(self, t_slot_total: int, pairleased_t: int = 0):
|
||||
"""检查 T slot 守恒。
|
||||
|
||||
Args:
|
||||
t_slot_total: T_Slot_Sem 初始容量
|
||||
pairleased_t: 当前 PairLease 持有的 T 数量
|
||||
"""
|
||||
free = self.t_slot_sem._value
|
||||
inv = self.inventory.t_depth
|
||||
held = self._worker_held_t
|
||||
total = free + inv + pairleased_t + held
|
||||
assert total == t_slot_total, (
|
||||
f'T slot 守恒违反: free({free}) + inv({inv}) + pairleased({pairleased_t}) '
|
||||
f'+ worker_held({held}) = {total} != total({t_slot_total})'
|
||||
)
|
||||
|
||||
def check_q_conservation(self, q_slot_total: int, pairleased_q: int = 0):
|
||||
"""检查 Q slot 守恒。"""
|
||||
free = self.q_slot_sem._value
|
||||
inv = self.inventory.q_depth
|
||||
held = self._worker_held_q
|
||||
total = free + inv + pairleased_q + held
|
||||
assert total == q_slot_total, (
|
||||
f'Q slot 守恒违反: free({free}) + inv({inv}) + pairleased({pairleased_q}) '
|
||||
f'+ worker_held({held}) = {total} != total({q_slot_total})'
|
||||
)
|
||||
|
||||
def check_pending_conservation(self, pending_total: int, in_flight: int = 0):
|
||||
"""检查 Q pending 守恒。"""
|
||||
free = self.q_pending_sem._value
|
||||
total = free + in_flight
|
||||
assert total == pending_total, (
|
||||
f'Q pending 守恒违反: free({free}) + in_flight({in_flight}) = {total} '
|
||||
f'!= total({pending_total})'
|
||||
)
|
||||
|
||||
def check_all(self, t_total, q_total, pend_total, pairleased_t=0, pairleased_q=0, in_flight=0):
|
||||
"""一次性检查所有守恒。"""
|
||||
self.check_t_conservation(t_total, pairleased_t)
|
||||
self.check_q_conservation(q_total, pairleased_q)
|
||||
self.check_pending_conservation(pend_total, in_flight)
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# 取消注入器
|
||||
# ═══════════════════════════════════════════
|
||||
class CancelInjector:
|
||||
"""在指定 await 边界前后注入取消。"""
|
||||
|
||||
def __init__(self):
|
||||
self._points: dict[str, asyncio.Event] = {}
|
||||
self._cancel_task: Optional[asyncio.Task] = None
|
||||
|
||||
def register(self, name: str):
|
||||
"""注册一个取消点。"""
|
||||
self._points[name] = asyncio.Event()
|
||||
|
||||
async def wait_at(self, name: str):
|
||||
"""Worker 在此挂起,等待外部触发取消。"""
|
||||
if name in self._points:
|
||||
await self._points[name].wait()
|
||||
|
||||
def trigger(self, name: str, task: asyncio.Task):
|
||||
"""触发指定取消点,取消目标 task。"""
|
||||
if name in self._points:
|
||||
self._points[name].set()
|
||||
task.cancel()
|
||||
|
||||
async def cancel_after(self, delay: float, task: asyncio.Task):
|
||||
"""延迟后取消目标 task。"""
|
||||
await asyncio.sleep(delay)
|
||||
task.cancel()
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# 事件记录器(用于性质测试)
|
||||
# ═══════════════════════════════════════════
|
||||
@dataclasses.dataclass
|
||||
class Event:
|
||||
ts: float
|
||||
kind: str # 't_produced', 't_admitted', 't_claimed', 'q_sent', 'q_returned',
|
||||
# 'q_admitted', 'q_claimed', 'pair_claimed', 'pair_ok', 'pair_fail',
|
||||
# 't_expired', 'q_expired', 't_discarded', 'q_discarded',
|
||||
# 'sem_t_free', 'sem_q_free', 'sem_pend_free', 'inv_t', 'inv_q'
|
||||
detail: Any = None
|
||||
|
||||
|
||||
class EventLog:
|
||||
"""记录所有事件,用于事后分析。"""
|
||||
|
||||
def __init__(self):
|
||||
self.events: list[Event] = []
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def record(self, kind: str, detail=None):
|
||||
async with self._lock:
|
||||
self.events.append(Event(ts=time.time(), kind=kind, detail=detail))
|
||||
|
||||
def count(self, kind: str) -> int:
|
||||
return sum(1 for e in self.events if e.kind == kind)
|
||||
|
||||
def filter(self, kind: str) -> list[Event]:
|
||||
return [e for e in self.events if e.kind == kind]
|
||||
|
||||
def dump(self) -> str:
|
||||
lines = []
|
||||
for e in self.events:
|
||||
lines.append(f'{e.ts:.3f} {e.kind} {e.detail or ""}')
|
||||
return '\n'.join(lines)
|
||||
2
tests/requirements.txt
Normal file
2
tests/requirements.txt
Normal file
@@ -0,0 +1,2 @@
|
||||
pytest>=8.0.0
|
||||
pytest-asyncio>=0.24.0
|
||||
82
tests/test_admission_gate.py
Normal file
82
tests/test_admission_gate.py
Normal file
@@ -0,0 +1,82 @@
|
||||
import asyncio
|
||||
import unittest
|
||||
|
||||
from grok_register.core.admission import AdmissionGate
|
||||
|
||||
|
||||
class FakeInventory:
|
||||
def __init__(self):
|
||||
self.t_depth = 0
|
||||
self.q_depth = 0
|
||||
|
||||
|
||||
class AdmissionGateTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_t_production_blocks_at_high_water_and_wakes_at_low_water(self):
|
||||
inventory = FakeInventory()
|
||||
gate = AdmissionGate(inventory, t_low=1, t_high=3, q_low=1, q_high=3)
|
||||
|
||||
inventory.t_depth = 3
|
||||
blocked = asyncio.create_task(gate.acquire_t_production())
|
||||
await asyncio.sleep(0.02)
|
||||
self.assertFalse(blocked.done())
|
||||
|
||||
inventory.t_depth = 1
|
||||
await gate.notify_changed()
|
||||
lease = await asyncio.wait_for(blocked, timeout=1)
|
||||
|
||||
self.assertEqual(gate.t_in_progress, 1)
|
||||
await lease.release()
|
||||
self.assertEqual(gate.t_in_progress, 0)
|
||||
|
||||
async def test_q_batch_reserves_only_current_demand(self):
|
||||
inventory = FakeInventory()
|
||||
gate = AdmissionGate(inventory, t_low=1, t_high=3, q_low=2, q_high=5)
|
||||
|
||||
lease = await gate.acquire_q_batch(max_batch=4)
|
||||
|
||||
self.assertEqual(lease.count, 4)
|
||||
self.assertEqual(gate.q_inflight, 4)
|
||||
|
||||
second = await gate.acquire_q_batch(max_batch=4)
|
||||
self.assertEqual(second.count, 1)
|
||||
self.assertEqual(gate.q_inflight, 5)
|
||||
|
||||
blocked = asyncio.create_task(gate.acquire_q_batch(max_batch=4))
|
||||
await asyncio.sleep(0.02)
|
||||
self.assertFalse(blocked.done())
|
||||
|
||||
await lease.release_all()
|
||||
await second.release_all()
|
||||
inventory.q_depth = 1
|
||||
await gate.notify_changed()
|
||||
third = await asyncio.wait_for(blocked, timeout=1)
|
||||
self.assertEqual(third.count, 4)
|
||||
await third.release_all()
|
||||
|
||||
async def test_q_batch_waits_for_low_water_after_reaching_high_water(self):
|
||||
inventory = FakeInventory()
|
||||
gate = AdmissionGate(inventory, t_low=1, t_high=3, q_low=2, q_high=5)
|
||||
|
||||
lease = await gate.acquire_q_batch(max_batch=5)
|
||||
self.assertEqual(lease.count, 5)
|
||||
self.assertEqual(gate.q_inflight, 5)
|
||||
|
||||
await lease.release_one()
|
||||
self.assertEqual(gate.q_inflight, 4)
|
||||
|
||||
blocked = asyncio.create_task(gate.acquire_q_batch(max_batch=5))
|
||||
await asyncio.sleep(0.02)
|
||||
self.assertFalse(blocked.done())
|
||||
|
||||
await lease.release_one()
|
||||
await lease.release_one()
|
||||
await lease.release_one()
|
||||
|
||||
resumed = await asyncio.wait_for(blocked, timeout=1)
|
||||
self.assertEqual(resumed.count, 4)
|
||||
await lease.release_all()
|
||||
await resumed.release_all()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
475
tests/test_cancel.py
Normal file
475
tests/test_cancel.py
Normal file
@@ -0,0 +1,475 @@
|
||||
"""
|
||||
Layer 2: 取消注入测试
|
||||
|
||||
在每个 await 边界前后强制取消 Worker,验证:
|
||||
- S 取消后 Physical_Sem 和 T_Slot_Sem 不泄漏
|
||||
- P 取消后 Physical_Sem、Q_Pending_Sem、Q_Slot_Sem 不泄漏
|
||||
- C 取消后 Physical_Sem 不泄漏,已 claim 的 T/Q 被核销
|
||||
- 库存守恒: 取消不导致 slot 无界堆积或耗尽
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
import pytest
|
||||
|
||||
from grok_register.core.envelope import ResourceEnvelope
|
||||
from grok_register.core.inventory import Inventory
|
||||
from grok_register.core.observer import Metrics
|
||||
from tests.fakes import FakeTurnstile, FakeEmailService, FakeRegisterAPI, Conservation
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# S_Worker 取消测试
|
||||
# ═══════════════════════════════════════════
|
||||
|
||||
async def _s_worker_impl(browser, inventory, physical_sem, t_slot_sem, metrics, stop):
|
||||
"""S_Worker 简化实现,用于测试。"""
|
||||
while not stop.is_set():
|
||||
await physical_sem.acquire()
|
||||
try:
|
||||
token = await browser.generate()
|
||||
finally:
|
||||
physical_sem.release()
|
||||
|
||||
if token is None:
|
||||
metrics.t_discarded += 1
|
||||
await asyncio.sleep(0.05)
|
||||
continue
|
||||
|
||||
metrics.t_produced += 1
|
||||
now = time.time()
|
||||
try:
|
||||
t_env = await ResourceEnvelope.create_with_slot(
|
||||
'T', token, t_slot_sem, expires_at=now + 300
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
metrics.t_discarded += 1
|
||||
raise
|
||||
|
||||
try:
|
||||
await inventory.put_t(t_env)
|
||||
except Exception:
|
||||
t_env.discard(reason='put_failed')
|
||||
metrics.t_discarded += 1
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
|
||||
class TestSCancel:
|
||||
"""S_Worker 取消语义测试。"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_during_generate_no_leak(self):
|
||||
"""S 在 generate() 期间被取消,sem 不泄漏。"""
|
||||
ft = FakeTurnstile(delay=0.5) # 慢生成
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
t_slot = asyncio.Semaphore(4)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
|
||||
task = asyncio.create_task(
|
||||
_s_worker_impl(ft, inv, phys, t_slot, metrics, stop)
|
||||
)
|
||||
await asyncio.sleep(0.05) # 等它进入 generate
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# Physical_Sem 应完全释放
|
||||
assert phys._value == 2, f'Physical_Sem 泄漏: {phys._value}'
|
||||
# T_Slot_Sem 不变(没生成成功)
|
||||
assert t_slot._value == 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_during_put_t_no_leak(self):
|
||||
"""S 在 put_t 期间被取消(极端: put 内部取消),envelope 被 discard。"""
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
t_slot = asyncio.Semaphore(4)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
|
||||
# 快速生成
|
||||
ft = FakeTurnstile(delay=0.001)
|
||||
|
||||
task = asyncio.create_task(
|
||||
_s_worker_impl(ft, inv, phys, t_slot, metrics, stop)
|
||||
)
|
||||
# 在 put_t 之前的极短窗口取消
|
||||
await asyncio.sleep(0.002)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert phys._value == 2
|
||||
# t_slot 可能是 4 或 3,但不应是负数
|
||||
assert t_slot._value >= 0
|
||||
assert t_slot._value <= 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s_repeated_cancel_no_accumulation(self):
|
||||
"""反复取消 S 多次,sem 计数不累积。"""
|
||||
ft = FakeTurnstile(delay=0.01)
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
t_slot = asyncio.Semaphore(4)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
|
||||
for _ in range(20):
|
||||
task = asyncio.create_task(
|
||||
_s_worker_impl(ft, inv, phys, t_slot, metrics, stop)
|
||||
)
|
||||
await asyncio.sleep(0.02)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert phys._value == 2, f'Physical_Sem 泄漏: {phys._value}'
|
||||
# 守恒: free + inventory == total (成功入库的 envelope 占着 slot)
|
||||
assert t_slot._value + inv.t_depth == 4, f'T_Slot 守恒违反: free={t_slot._value} inv={inv.t_depth}'
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# P_Worker 取消测试
|
||||
# ═══════════════════════════════════════════
|
||||
|
||||
async def _p_worker_impl(email_svc, inventory, physical_sem, q_pending_sem,
|
||||
q_slot_sem, metrics, stop):
|
||||
"""P_Worker 简化实现,用于测试。"""
|
||||
loop = asyncio.get_event_loop()
|
||||
while not stop.is_set():
|
||||
await q_pending_sem.acquire()
|
||||
_pending_released = False
|
||||
try:
|
||||
await physical_sem.acquire()
|
||||
try:
|
||||
handle, email, password = await email_svc.create_email()
|
||||
sent = True # fake 总是成功
|
||||
finally:
|
||||
physical_sem.release()
|
||||
|
||||
if not sent:
|
||||
q_pending_sem.release()
|
||||
_pending_released = True
|
||||
continue
|
||||
|
||||
metrics.q_sent += 1
|
||||
|
||||
try:
|
||||
code = await asyncio.wait_for(
|
||||
email_svc.poll_code(handle), timeout=95
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
code = None
|
||||
|
||||
if code is None:
|
||||
metrics.q_discarded += 1
|
||||
q_pending_sem.release()
|
||||
_pending_released = True
|
||||
continue
|
||||
|
||||
metrics.q_returned += 1
|
||||
now = time.time()
|
||||
q_env = await ResourceEnvelope.create_with_slot(
|
||||
'Q', {'email': email, 'password': password, 'code': code},
|
||||
q_slot_sem, expires_at=now + 120,
|
||||
)
|
||||
|
||||
try:
|
||||
await inventory.put_q(q_env)
|
||||
except Exception:
|
||||
q_env.discard(reason='put_failed')
|
||||
metrics.q_discarded += 1
|
||||
|
||||
except asyncio.CancelledError:
|
||||
if not _pending_released:
|
||||
q_pending_sem.release()
|
||||
_pending_released = True
|
||||
raise
|
||||
except Exception:
|
||||
metrics.q_discarded += 1
|
||||
finally:
|
||||
if not _pending_released:
|
||||
q_pending_sem.release()
|
||||
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
|
||||
class TestPCancel:
|
||||
"""P_Worker 取消语义测试。"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_during_create_email(self):
|
||||
"""P 在 create_email 期间被取消,所有 sem 正确释放。"""
|
||||
es = FakeEmailService()
|
||||
es.set_cancel_create_at(1)
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
q_pend = asyncio.Semaphore(4)
|
||||
q_slot = asyncio.Semaphore(4)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
|
||||
task = asyncio.create_task(
|
||||
_p_worker_impl(es, inv, phys, q_pend, q_slot, metrics, stop)
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert phys._value == 2, f'Physical_Sem 泄漏: {phys._value}'
|
||||
assert q_pend._value == 4, f'Q_Pending 泄漏: {q_pend._value}'
|
||||
assert q_slot._value == 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_during_poll_code(self):
|
||||
"""P 在 poll_code 期间被取消,不持有 Physical_Sem。"""
|
||||
es = FakeEmailService()
|
||||
es.set_cancel_poll_at(1)
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
q_pend = asyncio.Semaphore(4)
|
||||
q_slot = asyncio.Semaphore(4)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
|
||||
task = asyncio.create_task(
|
||||
_p_worker_impl(es, inv, phys, q_pend, q_slot, metrics, stop)
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# P 等 Q 时不持有 Physical_Sem
|
||||
assert phys._value == 2, f'Physical_Sem 应完全释放: {phys._value}'
|
||||
assert q_pend._value == 4, f'Q_Pending 泄漏: {q_pend._value}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_p_repeated_cancel_no_accumulation(self):
|
||||
"""反复取消 P 多次,sem 计数不累积。"""
|
||||
es = FakeEmailService()
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
q_pend = asyncio.Semaphore(4)
|
||||
q_slot = asyncio.Semaphore(4)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
|
||||
for _ in range(20):
|
||||
task = asyncio.create_task(
|
||||
_p_worker_impl(es, inv, phys, q_pend, q_slot, metrics, stop)
|
||||
)
|
||||
await asyncio.sleep(0.03)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert phys._value == 2, f'Physical_Sem 泄漏: {phys._value}'
|
||||
assert q_pend._value == 4, f'Q_Pending 泄漏: {q_pend._value}'
|
||||
# 守恒: free + inventory == total
|
||||
assert q_slot._value + inv.q_depth == 4, f'Q_Slot 守恒违反: free={q_slot._value} inv={inv.q_depth}'
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# C_Worker 取消测试
|
||||
# ═══════════════════════════════════════════
|
||||
|
||||
async def _c_worker_impl(register_api, browser_unused, inventory, physical_sem,
|
||||
metrics, stop, file_lock):
|
||||
"""C_Worker 简化实现,用于测试。"""
|
||||
while not stop.is_set():
|
||||
try:
|
||||
async with inventory.claim_pair() as pair:
|
||||
t_val = pair.t.value
|
||||
q_val = pair.q.value
|
||||
await physical_sem.acquire()
|
||||
try:
|
||||
sso = await register_api.register(
|
||||
q_val['email'], q_val['password'], q_val['code'], t_val
|
||||
)
|
||||
if sso:
|
||||
metrics.pair_consumed_ok += 1
|
||||
metrics.success_count += 1
|
||||
else:
|
||||
metrics.pair_consumed_fail += 1
|
||||
finally:
|
||||
physical_sem.release()
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
metrics.pair_consumed_fail += 1
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
|
||||
class TestCCancel:
|
||||
"""C_Worker 取消语义测试。"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_during_register_pair_settled(self):
|
||||
"""C 在 register 期间被取消,pair 已 claim,必须核销。"""
|
||||
fra = FakeRegisterAPI(delay=0.5)
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
t_slot = asyncio.Semaphore(4)
|
||||
q_slot = asyncio.Semaphore(4)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
fl = asyncio.Lock()
|
||||
|
||||
# 用带 metrics 的 inventory
|
||||
inv_with_metrics = Inventory(metrics=Metrics())
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'tok', t_slot)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'email': 'a@b.com', 'password': 'p', 'code': '123'}, q_slot)
|
||||
await inv_with_metrics.put_t(t_env)
|
||||
await inv_with_metrics.put_q(q_env)
|
||||
|
||||
task = asyncio.create_task(
|
||||
_c_worker_impl(fra, None, inv_with_metrics, phys, metrics, stop, fl)
|
||||
)
|
||||
await asyncio.sleep(0.1) # 等它 claim 进入 register
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# pair 已核销: slot 应释放
|
||||
assert t_slot._value == 4, f'T slot 未核销: {t_slot._value}'
|
||||
assert q_slot._value == 4, f'Q slot 未核销: {q_slot._value}'
|
||||
assert phys._value == 2, f'Physical_Sem 泄漏: {phys._value}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_during_claim_wait_no_pop(self):
|
||||
"""C 在 claim_pair 等待时取消,不弹出资源。"""
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
t_slot = asyncio.Semaphore(4)
|
||||
q_slot = asyncio.Semaphore(4)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
fl = asyncio.Lock()
|
||||
|
||||
fra = FakeRegisterAPI()
|
||||
|
||||
# 放一个 T 但不放 Q,让 C 等待
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'tok', t_slot)
|
||||
await inv.put_t(t_env)
|
||||
|
||||
task = asyncio.create_task(
|
||||
_c_worker_impl(fra, None, inv, phys, metrics, stop, fl)
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# T 仍在 inventory 中
|
||||
assert inv.t_depth == 1, 'claim 等待时取消不应弹出 T'
|
||||
assert phys._value == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_c_repeated_cancel_no_accumulation(self):
|
||||
"""反复取消 C 多次,sem 计数不累积。"""
|
||||
fra = FakeRegisterAPI(delay=0.01)
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
t_slot = asyncio.Semaphore(10)
|
||||
q_slot = asyncio.Semaphore(10)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
fl = asyncio.Lock()
|
||||
|
||||
for _ in range(10):
|
||||
# 每次放一对
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', f'tok_{_}', t_slot)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'email': f'e@{_}', 'password': 'p', 'code': '123'}, q_slot)
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
task = asyncio.create_task(
|
||||
_c_worker_impl(fra, None, inv, phys, metrics, stop, fl)
|
||||
)
|
||||
await asyncio.sleep(0.03)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert phys._value == 2, f'Physical_Sem 泄漏: {phys._value}'
|
||||
# 所有 slot 应已核销
|
||||
assert t_slot._value == 10, f'T_Slot 泄漏: {t_slot._value}'
|
||||
assert q_slot._value == 10, f'Q_Slot 泄漏: {q_slot._value}'
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# 端到端取消: S+P+C 同时运行时取消
|
||||
# ═══════════════════════════════════════════
|
||||
class TestE2ECancel:
|
||||
"""端到端取消测试: S/P/C 同时运行时取消,验证全局守恒。"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_all_workers_sem_reset(self):
|
||||
"""同时启动 S/P/C,然后全部取消,所有 sem 恢复初始值。"""
|
||||
ft = FakeTurnstile(delay=0.05)
|
||||
es = FakeEmailService()
|
||||
fra = FakeRegisterAPI(delay=0.05)
|
||||
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(4)
|
||||
t_slot = asyncio.Semaphore(8)
|
||||
q_slot = asyncio.Semaphore(8)
|
||||
q_pend = asyncio.Semaphore(12)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
fl = asyncio.Lock()
|
||||
|
||||
tasks = []
|
||||
for _ in range(3):
|
||||
tasks.append(asyncio.create_task(
|
||||
_s_worker_impl(ft, inv, phys, t_slot, metrics, stop)
|
||||
))
|
||||
tasks.append(asyncio.create_task(
|
||||
_p_worker_impl(es, inv, phys, q_pend, q_slot, metrics, stop)
|
||||
))
|
||||
tasks.append(asyncio.create_task(
|
||||
_c_worker_impl(fra, None, inv, phys, metrics, stop, fl)
|
||||
))
|
||||
|
||||
await asyncio.sleep(0.2)
|
||||
stop.set()
|
||||
for t in tasks:
|
||||
t.cancel()
|
||||
for t in tasks:
|
||||
try:
|
||||
await t
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# 全部停止后,守恒检查: free + inventory == total
|
||||
assert phys._value >= 0 and phys._value <= 4
|
||||
assert t_slot._value + inv.t_depth == 8, f'T 守恒违反: free={t_slot._value} inv={inv.t_depth}'
|
||||
assert q_slot._value + inv.q_depth == 8, f'Q 守恒违反: free={q_slot._value} inv={inv.q_depth}'
|
||||
assert q_pend._value >= 0 and q_pend._value <= 12
|
||||
|
||||
# 守恒检查
|
||||
cons = Conservation(t_slot, q_slot, q_pend, inv)
|
||||
cons.check_t_conservation(8)
|
||||
cons.check_q_conservation(8)
|
||||
cons.check_pending_conservation(12)
|
||||
491
tests/test_core.py
Normal file
491
tests/test_core.py
Normal file
@@ -0,0 +1,491 @@
|
||||
"""
|
||||
Layer 1: 单元测试
|
||||
|
||||
测 ResourceEnvelope、PairLease、Inventory 的局部语义:
|
||||
- release_slot_once() 幂等
|
||||
- create_with_slot() slot 获取失败不创建 envelope
|
||||
- is_expired() 正确判断
|
||||
- put_t/put_q 成功前所有权属于 Worker,成功后属于 Inventory
|
||||
- claim_pair 等待时取消不弹资源
|
||||
- claim_pair 成功后 pair 属于 PairLease
|
||||
- PairLease 退出时核销 T/Q
|
||||
- 过期资源不会交付给 C
|
||||
- lazy cleanup 只在事件触发时发生
|
||||
- discard() 幂等释放
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
import pytest
|
||||
|
||||
from grok_register.core.envelope import ResourceEnvelope
|
||||
from grok_register.core.inventory import Inventory, PairLease
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# ResourceEnvelope 测试
|
||||
# ═══════════════════════════════════════════
|
||||
class TestResourceEnvelope:
|
||||
"""ResourceEnvelope 语义测试。"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_with_slot_acquires_sem(self):
|
||||
"""create_with_slot 成功后 semaphore 值减 1。"""
|
||||
sem = asyncio.Semaphore(3)
|
||||
env = await ResourceEnvelope.create_with_slot('T', 'tok1', sem)
|
||||
assert sem._value == 2
|
||||
assert env.value == 'tok1'
|
||||
assert env.kind == 'T'
|
||||
assert not env.released
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_with_slot_fails_no_envelope(self):
|
||||
"""slot 获取失败(取消)时不创建 envelope,sem 不变。"""
|
||||
sem = asyncio.Semaphore(0)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
# 手动取消
|
||||
task = asyncio.current_task()
|
||||
task.cancel()
|
||||
try:
|
||||
await ResourceEnvelope.create_with_slot('T', 'tok1', sem)
|
||||
except asyncio.CancelledError:
|
||||
# sem 不应改变(因为 acquire 被取消了)
|
||||
raise
|
||||
# sem 仍为 0
|
||||
assert sem._value == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_slot_once_idempotent(self):
|
||||
"""release_slot_once() 幂等:重复调用只释放一次。"""
|
||||
sem = asyncio.Semaphore(3)
|
||||
env = await ResourceEnvelope.create_with_slot('T', 'tok1', sem)
|
||||
assert sem._value == 2
|
||||
|
||||
env.release_slot_once(reason='test1')
|
||||
assert sem._value == 3
|
||||
assert env.released
|
||||
|
||||
# 第二次调用不应再次释放
|
||||
env.release_slot_once(reason='test2')
|
||||
assert sem._value == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discard_releases_slot(self):
|
||||
"""discard() 释放 slot。"""
|
||||
sem = asyncio.Semaphore(2)
|
||||
env = await ResourceEnvelope.create_with_slot('Q', {'code': '123'}, sem)
|
||||
assert sem._value == 1
|
||||
env.discard(reason='test')
|
||||
assert sem._value == 2
|
||||
assert env.released
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discard_idempotent(self):
|
||||
"""discard() 幂等。"""
|
||||
sem = asyncio.Semaphore(2)
|
||||
env = await ResourceEnvelope.create_with_slot('Q', 'val', sem)
|
||||
env.discard(reason='first')
|
||||
env.discard(reason='second')
|
||||
assert sem._value == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_expired_with_expires_at(self):
|
||||
"""有 expires_at 时正确判断过期。"""
|
||||
sem = asyncio.Semaphore(2)
|
||||
now = time.time()
|
||||
env = await ResourceEnvelope.create_with_slot('T', 'tok', sem, expires_at=now - 1)
|
||||
assert env.is_expired(now=now)
|
||||
|
||||
env2 = await ResourceEnvelope.create_with_slot('T', 'tok2', sem, expires_at=now + 100)
|
||||
assert not env2.is_expired(now=now)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_expired_without_expires_at(self):
|
||||
"""没有 expires_at 时永远不过期。"""
|
||||
sem = asyncio.Semaphore(2)
|
||||
env = await ResourceEnvelope.create_with_slot('T', 'tok', sem)
|
||||
assert not env.is_expired(now=time.time() + 99999)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_with_slot_meta(self):
|
||||
"""meta 字段正确传递。"""
|
||||
sem = asyncio.Semaphore(1)
|
||||
env = await ResourceEnvelope.create_with_slot('T', 'tok', sem, meta={'source': 'test'})
|
||||
assert env.meta == {'source': 'test'}
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# Inventory 测试
|
||||
# ═══════════════════════════════════════════
|
||||
class TestInventory:
|
||||
"""Inventory 语义测试。"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_put_t_then_claim(self):
|
||||
"""put_t 后 claim_pair 能拿到 pair。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
q_sem = asyncio.Semaphore(2)
|
||||
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'token1', t_sem)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': 'abc'}, q_sem)
|
||||
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
assert inv.t_depth == 1
|
||||
assert inv.q_depth == 1
|
||||
|
||||
# claim
|
||||
async with inv.claim_pair() as pair:
|
||||
assert pair.t.value == 'token1'
|
||||
assert pair.q.value == {'code': 'abc'}
|
||||
# claim 后 inventory 应为空
|
||||
assert inv.t_depth == 0
|
||||
assert inv.q_depth == 0
|
||||
|
||||
# PairLease 退出后 slot 应被释放
|
||||
assert t_sem._value == 2
|
||||
assert q_sem._value == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claim_waits_for_both(self):
|
||||
"""只有 T 没有 Q 时,claim_pair 应等待。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
q_sem = asyncio.Semaphore(2)
|
||||
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'tok', t_sem)
|
||||
await inv.put_t(t_env)
|
||||
|
||||
claimed = False
|
||||
|
||||
async def claimer():
|
||||
nonlocal claimed
|
||||
async with inv.claim_pair() as pair:
|
||||
claimed = True
|
||||
|
||||
task = asyncio.create_task(claimer())
|
||||
await asyncio.sleep(0.1)
|
||||
assert not claimed, 'claim 不应在缺少 Q 时成功'
|
||||
|
||||
# 现在放入 Q
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': 'x'}, q_sem)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
await asyncio.wait_for(task, timeout=2)
|
||||
assert claimed
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claim_cancel_no_pop(self):
|
||||
"""claim_pair 等待时取消,不应弹出任何资源。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'tok', t_sem)
|
||||
await inv.put_t(t_env)
|
||||
assert inv.t_depth == 1
|
||||
|
||||
async def claimer():
|
||||
async with inv.claim_pair() as pair:
|
||||
pass # 不应到达
|
||||
|
||||
task = asyncio.create_task(claimer())
|
||||
await asyncio.sleep(0.1)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# T 仍在 inventory 中
|
||||
assert inv.t_depth == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pair_lease_settles_on_success(self):
|
||||
"""PairLease 成功退出时核销 T/Q 并释放 slot。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
q_sem = asyncio.Semaphore(2)
|
||||
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'tok', t_sem)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': 'x'}, q_sem)
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
async with inv.claim_pair() as pair:
|
||||
# slot 还在 pair 中
|
||||
assert t_sem._value == 1
|
||||
assert q_sem._value == 1
|
||||
|
||||
# 退出后 slot 释放
|
||||
assert t_sem._value == 2
|
||||
assert q_sem._value == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pair_lease_settles_on_exception(self):
|
||||
"""PairLease 异常退出时也核销 T/Q。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
q_sem = asyncio.Semaphore(2)
|
||||
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'tok', t_sem)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': 'x'}, q_sem)
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
async with inv.claim_pair() as pair:
|
||||
raise ValueError('boom')
|
||||
|
||||
# slot 仍被核销
|
||||
assert t_sem._value == 2
|
||||
assert q_sem._value == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pair_lease_settles_on_cancel(self):
|
||||
"""PairLease 取消退出时也核销 T/Q。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
q_sem = asyncio.Semaphore(2)
|
||||
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'tok', t_sem)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': 'x'}, q_sem)
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
async def consumer():
|
||||
async with inv.claim_pair() as pair:
|
||||
# 在 pair 内部被取消
|
||||
await asyncio.sleep(999)
|
||||
|
||||
task = asyncio.create_task(consumer())
|
||||
await asyncio.sleep(0.1)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# slot 仍被核销
|
||||
assert t_sem._value == 2
|
||||
assert q_sem._value == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_expired_t_not_delivered(self):
|
||||
"""过期的 T 不会交付给 C。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
q_sem = asyncio.Semaphore(2)
|
||||
|
||||
now = time.time()
|
||||
# T 已过期
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'expired_tok', t_sem, expires_at=now - 10)
|
||||
# Q 有效
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': 'x'}, q_sem, expires_at=now + 100)
|
||||
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
# claim 应该把过期的 T 清理掉,然后等待新的 T
|
||||
claimed = False
|
||||
|
||||
async def claimer():
|
||||
nonlocal claimed
|
||||
async with inv.claim_pair() as pair:
|
||||
claimed = True
|
||||
|
||||
task = asyncio.create_task(claimer())
|
||||
await asyncio.sleep(0.2)
|
||||
assert not claimed, '过期 T 不应交付'
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# 过期 T 被清理,slot 被释放
|
||||
assert t_sem._value == 2
|
||||
# Q 仍在 inventory 中
|
||||
assert inv.q_depth == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_expired_q_not_delivered(self):
|
||||
"""过期的 Q 不会交付给 C。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
q_sem = asyncio.Semaphore(2)
|
||||
|
||||
now = time.time()
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'tok', t_sem, expires_at=now + 100)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': 'x'}, q_sem, expires_at=now - 10)
|
||||
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
claimed = False
|
||||
|
||||
async def claimer():
|
||||
nonlocal claimed
|
||||
async with inv.claim_pair() as pair:
|
||||
claimed = True
|
||||
|
||||
task = asyncio.create_task(claimer())
|
||||
await asyncio.sleep(0.2)
|
||||
assert not claimed, '过期 Q 不应交付'
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert q_sem._value == 2
|
||||
assert inv.t_depth == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lazy_cleanup_on_put(self):
|
||||
"""lazy cleanup: put 时清理已过期的库存。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(4)
|
||||
|
||||
now = time.time()
|
||||
# 放入 3 个过期的 T
|
||||
for i in range(3):
|
||||
env = await ResourceEnvelope.create_with_slot(
|
||||
'T', f'expired_{i}', t_sem, expires_at=now - 10
|
||||
)
|
||||
await inv.put_t(env)
|
||||
|
||||
# 由于 put 时 lazy cleanup 已清理过期项,库存应为空
|
||||
assert inv.t_depth == 0
|
||||
# slot 应已释放
|
||||
assert t_sem._value == 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lazy_cleanup_silent_until_next_event(self):
|
||||
"""lazy cleanup 不在静默期主动释放,下一次 put 才触发清理。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(4)
|
||||
|
||||
env = await ResourceEnvelope.create_with_slot(
|
||||
'T', 'will_expire', t_sem, expires_at=time.time() + 0.05
|
||||
)
|
||||
await inv.put_t(env)
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
assert inv.t_depth == 1
|
||||
assert t_sem._value == 3
|
||||
|
||||
fresh = await ResourceEnvelope.create_with_slot(
|
||||
'T', 'fresh', t_sem, expires_at=time.time() + 100
|
||||
)
|
||||
await inv.put_t(fresh)
|
||||
|
||||
assert inv.t_depth == 1
|
||||
assert t_sem._value == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lazy_cleanup_on_claim(self):
|
||||
"""lazy cleanup: claim_pair 时清理过期库存。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(4)
|
||||
q_sem = asyncio.Semaphore(4)
|
||||
|
||||
now = time.time()
|
||||
# 直接注入过期项到 inventory 内部(绕过 put 的清理)
|
||||
env = await ResourceEnvelope.create_with_slot('T', 'expired', t_sem, expires_at=now - 10)
|
||||
inv._t_buf.append(env)
|
||||
|
||||
# 有效的 Q
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': 'x'}, q_sem, expires_at=now + 100)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
# claim 时应清理过期的 T
|
||||
claimed = False
|
||||
async def claimer():
|
||||
nonlocal claimed
|
||||
async with inv.claim_pair() as pair:
|
||||
claimed = True
|
||||
|
||||
task = asyncio.create_task(claimer())
|
||||
await asyncio.sleep(0.2)
|
||||
assert not claimed # T 被清了,没有 pair
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert inv.t_depth == 0
|
||||
assert t_sem._value == 4 # slot 已释放
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fifo_order(self):
|
||||
"""Inventory 按 FIFO 顺序配对。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(10)
|
||||
q_sem = asyncio.Semaphore(10)
|
||||
|
||||
for i in range(3):
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', f'tok_{i}', t_sem)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': f'c_{i}'}, q_sem)
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
pairs = []
|
||||
for _ in range(3):
|
||||
async with inv.claim_pair() as pair:
|
||||
pairs.append((pair.t.value, pair.q.value['code']))
|
||||
|
||||
assert pairs == [('tok_0', 'c_0'), ('tok_1', 'c_1'), ('tok_2', 'c_2')]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_concurrent_claims(self):
|
||||
"""多个 C 同时 claim,不应重复消费同一个 T/Q。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(10)
|
||||
q_sem = asyncio.Semaphore(10)
|
||||
|
||||
n = 5
|
||||
for i in range(n):
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', f'tok_{i}', t_sem)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': f'c_{i}'}, q_sem)
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
results = []
|
||||
|
||||
async def claimer(idx):
|
||||
async with inv.claim_pair() as pair:
|
||||
results.append((idx, pair.t.value))
|
||||
|
||||
tasks = [asyncio.create_task(claimer(i)) for i in range(n)]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
# 每个 T 只被消费一次
|
||||
t_values = [r[1] for r in results]
|
||||
assert len(set(t_values)) == n, f'重复消费: {t_values}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metrics_admitted(self):
|
||||
"""put 成功时 metrics.t_admitted / q_admitted 递增。"""
|
||||
from grok_register.core.observer import Metrics
|
||||
m = Metrics()
|
||||
inv = Inventory(metrics=m)
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
|
||||
env = await ResourceEnvelope.create_with_slot('T', 'tok', t_sem)
|
||||
await inv.put_t(env)
|
||||
assert m.t_admitted == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metrics_expired(self):
|
||||
"""过期资源 put 时 metrics.t_expired 递增。"""
|
||||
from grok_register.core.observer import Metrics
|
||||
m = Metrics()
|
||||
inv = Inventory(metrics=m)
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
|
||||
env = await ResourceEnvelope.create_with_slot('T', 'tok', t_sem, expires_at=time.time() - 10)
|
||||
await inv.put_t(env)
|
||||
assert m.t_expired == 1
|
||||
assert m.t_admitted == 0
|
||||
58
tests/test_inventory_unittest.py
Normal file
58
tests/test_inventory_unittest.py
Normal file
@@ -0,0 +1,58 @@
|
||||
import asyncio
|
||||
import time
|
||||
import unittest
|
||||
|
||||
from grok_register.core.envelope import ResourceEnvelope
|
||||
from grok_register.core.inventory import Inventory
|
||||
from grok_register.core.observer import Metrics
|
||||
|
||||
|
||||
class InventoryUnitTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_create_with_slot_releases_slot_if_constructor_fails(self):
|
||||
class BrokenEnvelope(ResourceEnvelope):
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
sem = asyncio.Semaphore(1)
|
||||
|
||||
with self.assertRaises(RuntimeError):
|
||||
await BrokenEnvelope.create_with_slot("T", "tok", sem)
|
||||
|
||||
self.assertEqual(sem._value, 1)
|
||||
|
||||
async def test_lazy_cleanup_scans_past_fresh_head(self):
|
||||
metrics = Metrics()
|
||||
inventory = Inventory(metrics=metrics)
|
||||
t_sem = asyncio.Semaphore(3)
|
||||
q_sem = asyncio.Semaphore(1)
|
||||
now = time.time()
|
||||
|
||||
fresh_head = await ResourceEnvelope.create_with_slot(
|
||||
"T", "fresh_head", t_sem, expires_at=now + 100
|
||||
)
|
||||
stale_tail = await ResourceEnvelope.create_with_slot(
|
||||
"T", "stale_tail", t_sem, expires_at=now + 100
|
||||
)
|
||||
fresh_tail = await ResourceEnvelope.create_with_slot(
|
||||
"T", "fresh_tail", t_sem, expires_at=now + 100
|
||||
)
|
||||
|
||||
await inventory.put_t(fresh_head)
|
||||
await inventory.put_t(stale_tail)
|
||||
await inventory.put_t(fresh_tail)
|
||||
self.assertEqual(t_sem._value, 0)
|
||||
|
||||
stale_tail.expires_at = time.time() - 10
|
||||
|
||||
q_env = await ResourceEnvelope.create_with_slot(
|
||||
"Q", "q", q_sem, expires_at=now + 100
|
||||
)
|
||||
await inventory.put_q(q_env)
|
||||
|
||||
self.assertEqual(t_sem._value, 1)
|
||||
self.assertEqual(metrics.t_expired, 1)
|
||||
self.assertEqual(inventory.t_depth, 2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
482
tests/test_property.py
Normal file
482
tests/test_property.py
Normal file
@@ -0,0 +1,482 @@
|
||||
"""
|
||||
Layer 3: 性质测试 (Property Tests)
|
||||
|
||||
随机生成 S/P/C 成功、失败、取消、超时、过期事件,持续检查:
|
||||
- slot 守恒 (T/Q 总量恒定)
|
||||
- 无重复 claim (同一个 T/Q 不被两个 PairLease 拿到)
|
||||
- 无单边持有等待 (C 不先拿 T 再等 Q)
|
||||
- 过期资源不交付
|
||||
- 取消后不泄漏
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
import random
|
||||
import pytest
|
||||
|
||||
from grok_register.core.envelope import ResourceEnvelope
|
||||
from grok_register.core.inventory import Inventory
|
||||
from grok_register.core.observer import Metrics
|
||||
from tests.fakes import FakeTurnstile, FakeEmailService, FakeRegisterAPI, Conservation, EventLog
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# 随机性质测试引擎
|
||||
# ═══════════════════════════════════════════
|
||||
|
||||
class PropertyEngine:
|
||||
"""驱动 S/P/C 随机行为,持续检查不变量。"""
|
||||
|
||||
def __init__(self, t_slot_cap=6, q_slot_cap=6, q_pending_cap=8, phys_cap=4,
|
||||
t_max_age=2.0, q_max_age=2.0, seed=None):
|
||||
self.rng = random.Random(seed)
|
||||
self.t_slot_sem = asyncio.Semaphore(t_slot_cap)
|
||||
self.q_slot_sem = asyncio.Semaphore(q_slot_cap)
|
||||
self.q_pending_sem = asyncio.Semaphore(q_pending_cap)
|
||||
self.phys_sem = asyncio.Semaphore(phys_cap)
|
||||
self.t_slot_cap = t_slot_cap
|
||||
self.q_slot_cap = q_slot_cap
|
||||
self.q_pending_cap = q_pending_cap
|
||||
self.t_max_age = t_max_age
|
||||
self.q_max_age = q_max_age
|
||||
|
||||
self.metrics = Metrics()
|
||||
self.inventory = Inventory(metrics=self.metrics)
|
||||
self.conservation = Conservation(
|
||||
self.t_slot_sem, self.q_slot_sem, self.q_pending_sem, self.inventory
|
||||
)
|
||||
self.events = EventLog()
|
||||
|
||||
self._claimed_tokens = set() # 检测重复 claim
|
||||
self._pairleased_t = 0 # 当前 PairLease 持有的 T
|
||||
self._pairleased_q = 0 # 当前 PairLease 持有的 Q
|
||||
|
||||
def check_invariants_soft(self):
|
||||
"""运行时软检查: 非负、不超限。
|
||||
|
||||
注意: Worker 运行期间 slot 可能在 create_with_slot 和 put_t 之间的
|
||||
envelope 中(worker 持有但无法从外部追踪)。严格守恒检查用
|
||||
check_invariants_final() 在所有 Worker 停止后执行。
|
||||
"""
|
||||
assert self.t_slot_sem._value >= 0, 'T slot 为负'
|
||||
assert self.q_slot_sem._value >= 0, 'Q slot 为负'
|
||||
assert self.q_pending_sem._value >= 0, 'Q pending 为负'
|
||||
assert self.q_pending_sem._value <= self.q_pending_cap, 'Q pending 超限'
|
||||
assert self.phys_sem._value >= 0, 'Physical 为负'
|
||||
# inventory 深度不超过 slot 总量
|
||||
assert self.inventory.t_depth <= self.t_slot_cap, 'T inventory 超限'
|
||||
assert self.inventory.q_depth <= self.q_slot_cap, 'Q inventory 超限'
|
||||
|
||||
def check_invariants_final(self):
|
||||
"""最终严格守恒检查: 所有 Worker 停止后执行。
|
||||
|
||||
此时没有 Worker 持有中间状态的 slot,也没有 PairLease。
|
||||
free + inventory == total
|
||||
"""
|
||||
self.conservation.check_t_conservation(self.t_slot_cap)
|
||||
self.conservation.check_q_conservation(self.q_slot_cap)
|
||||
self.conservation.check_pending_conservation(self.q_pending_cap)
|
||||
|
||||
async def s_produce(self):
|
||||
"""S 产生一个 T。返回 True/False/None(cancelled)。"""
|
||||
# 模拟物理资源使用
|
||||
await self.phys_sem.acquire()
|
||||
delay = self.rng.uniform(0.001, 0.02)
|
||||
try:
|
||||
await asyncio.sleep(delay)
|
||||
finally:
|
||||
self.phys_sem.release()
|
||||
|
||||
# 模拟随机失败
|
||||
if self.rng.random() < 0.1:
|
||||
self.metrics.t_discarded += 1
|
||||
await self.events.record('t_discarded', 'gen_fail')
|
||||
return False
|
||||
|
||||
token = f'tok_{self.metrics.t_produced}'
|
||||
self.metrics.t_produced += 1
|
||||
await self.events.record('t_produced', token)
|
||||
|
||||
# 创建 envelope + 入库
|
||||
# 关键: create_with_slot 成功后 slot 在 envelope 中,
|
||||
# 必须追踪 worker_held 直到 put_t 成功或 discard。
|
||||
now = time.time()
|
||||
expire = now + self.t_max_age if self.rng.random() > 0.15 else now - 1
|
||||
t_env = None
|
||||
try:
|
||||
t_env = await ResourceEnvelope.create_with_slot(
|
||||
'T', token, self.t_slot_sem, expires_at=expire
|
||||
)
|
||||
# slot 已获取,标记 worker 持有
|
||||
self.conservation.worker_acquire_t()
|
||||
await self.inventory.put_t(t_env)
|
||||
# put 成功,所有权转移到 inventory
|
||||
self.conservation.worker_release_t()
|
||||
except Exception:
|
||||
if t_env is not None and not t_env.released:
|
||||
t_env.discard(reason='error_cleanup')
|
||||
self.conservation.worker_release_t()
|
||||
self.metrics.t_discarded += 1
|
||||
return False
|
||||
except asyncio.CancelledError:
|
||||
if t_env is not None and not t_env.released:
|
||||
t_env.discard(reason='cancel_cleanup')
|
||||
self.conservation.worker_release_t()
|
||||
self.metrics.t_discarded += 1
|
||||
raise
|
||||
|
||||
self.check_invariants_soft()
|
||||
return True
|
||||
|
||||
async def p_produce(self):
|
||||
"""P 产生一个 Q。返回 True/False/None(cancelled)。"""
|
||||
await self.q_pending_sem.acquire()
|
||||
_pend_released = False
|
||||
|
||||
try:
|
||||
# 创建邮箱 + 发请求(持有 Physical)
|
||||
await self.phys_sem.acquire()
|
||||
try:
|
||||
delay = self.rng.uniform(0.001, 0.02)
|
||||
await asyncio.sleep(delay)
|
||||
finally:
|
||||
self.phys_sem.release()
|
||||
|
||||
# 模拟随机失败
|
||||
if self.rng.random() < 0.1:
|
||||
self.metrics.q_discarded += 1
|
||||
self.q_pending_sem.release()
|
||||
_pend_released = True
|
||||
await self.events.record('q_discarded', 'send_fail')
|
||||
return False
|
||||
|
||||
self.metrics.q_sent += 1
|
||||
|
||||
# 等待 Q 返回(不持有 Physical)
|
||||
poll_delay = self.rng.uniform(0.001, 0.05)
|
||||
await asyncio.sleep(poll_delay)
|
||||
|
||||
# 模拟超时/失败
|
||||
if self.rng.random() < 0.1:
|
||||
self.metrics.q_discarded += 1
|
||||
self.q_pending_sem.release()
|
||||
_pend_released = True
|
||||
await self.events.record('q_discarded', 'poll_timeout')
|
||||
return False
|
||||
|
||||
code = f'{self.rng.randint(100000, 999999)}'
|
||||
self.metrics.q_returned += 1
|
||||
await self.events.record('q_returned', code)
|
||||
|
||||
# 创建 envelope + 入库
|
||||
now = time.time()
|
||||
expire = now + self.q_max_age if self.rng.random() > 0.15 else now - 1
|
||||
q_env = None
|
||||
try:
|
||||
q_env = await ResourceEnvelope.create_with_slot(
|
||||
'Q', {'email': f'e@{code}', 'password': 'p', 'code': code},
|
||||
self.q_slot_sem, expires_at=expire,
|
||||
)
|
||||
self.conservation.worker_acquire_q()
|
||||
await self.inventory.put_q(q_env)
|
||||
self.conservation.worker_release_q()
|
||||
except Exception:
|
||||
if q_env is not None and not q_env.released:
|
||||
q_env.discard(reason='error_cleanup')
|
||||
self.conservation.worker_release_q()
|
||||
self.metrics.q_discarded += 1
|
||||
except asyncio.CancelledError:
|
||||
if q_env is not None and not q_env.released:
|
||||
q_env.discard(reason='cancel_cleanup')
|
||||
self.conservation.worker_release_q()
|
||||
self.metrics.q_discarded += 1
|
||||
raise
|
||||
|
||||
except asyncio.CancelledError:
|
||||
if not _pend_released:
|
||||
self.q_pending_sem.release()
|
||||
_pend_released = True
|
||||
raise
|
||||
except Exception:
|
||||
self.metrics.q_discarded += 1
|
||||
finally:
|
||||
if not _pend_released:
|
||||
self.q_pending_sem.release()
|
||||
|
||||
self.check_invariants_soft()
|
||||
return True
|
||||
|
||||
async def c_consume(self):
|
||||
"""C 消费一个 pair。返回 True/False/None(cancelled)。"""
|
||||
try:
|
||||
async with self.inventory.claim_pair() as pair:
|
||||
t_val = pair.t.value
|
||||
q_val = pair.q.value
|
||||
|
||||
# 检测重复 claim
|
||||
assert t_val not in self._claimed_tokens, f'重复 claim: {t_val}'
|
||||
self._claimed_tokens.add(t_val)
|
||||
|
||||
# pairleased 计数 — 必须在 claim 后立即标记,
|
||||
# 否则在下一个 await 前被取消会导致守恒违反
|
||||
self._pairleased_t += 1
|
||||
self._pairleased_q += 1
|
||||
|
||||
# 消费(持有 Physical)
|
||||
await self.phys_sem.acquire()
|
||||
try:
|
||||
delay = self.rng.uniform(0.001, 0.02)
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
# 模拟随机失败
|
||||
if self.rng.random() < 0.15:
|
||||
self.metrics.pair_consumed_fail += 1
|
||||
await self.events.record('pair_fail', t_val)
|
||||
else:
|
||||
self.metrics.pair_consumed_ok += 1
|
||||
self.metrics.success_count += 1
|
||||
await self.events.record('pair_ok', t_val)
|
||||
finally:
|
||||
self.phys_sem.release()
|
||||
|
||||
# PairLease 退出后(正常路径)
|
||||
self._pairleased_t -= 1
|
||||
self._pairleased_q -= 1
|
||||
|
||||
except asyncio.CancelledError:
|
||||
# PairLease.__aexit__ 已核销 slot。
|
||||
# 如果 _pairleased 已 +1,则 -1;否则不操作。
|
||||
if self._pairleased_t > 0:
|
||||
self._pairleased_t -= 1
|
||||
if self._pairleased_q > 0:
|
||||
self._pairleased_q -= 1
|
||||
raise
|
||||
except Exception:
|
||||
self.metrics.pair_consumed_fail += 1
|
||||
if self._pairleased_t > 0:
|
||||
self._pairleased_t -= 1
|
||||
if self._pairleased_q > 0:
|
||||
self._pairleased_q -= 1
|
||||
|
||||
self.check_invariants_soft()
|
||||
return True
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
class TestProperty:
|
||||
"""性质测试: 随机事件流 + 持续不变量检查。"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_random_events_10s(self):
|
||||
"""随机 S/P/C 事件流 10 秒,持续检查不变量。"""
|
||||
engine = PropertyEngine(seed=42)
|
||||
stop = asyncio.Event()
|
||||
errors = []
|
||||
|
||||
async def s_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.s_produce()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'S: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.005, 0.03))
|
||||
|
||||
async def p_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.p_produce()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'P: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.005, 0.03))
|
||||
|
||||
async def c_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.c_consume()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'C: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.005, 0.03))
|
||||
|
||||
tasks = []
|
||||
for _ in range(2):
|
||||
tasks.append(asyncio.create_task(s_loop()))
|
||||
tasks.append(asyncio.create_task(p_loop()))
|
||||
tasks.append(asyncio.create_task(c_loop()))
|
||||
|
||||
await asyncio.sleep(10)
|
||||
stop.set()
|
||||
for t in tasks:
|
||||
t.cancel()
|
||||
for t in tasks:
|
||||
try:
|
||||
await t
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# 最终不变量检查
|
||||
engine.check_invariants_final()
|
||||
|
||||
assert not errors, f'不变量违反:\n' + '\n'.join(errors)
|
||||
print(f'\n[Property] t_produced={engine.metrics.t_produced} '
|
||||
f'q_sent={engine.metrics.q_sent} q_returned={engine.metrics.q_returned} '
|
||||
f'pair_ok={engine.metrics.pair_consumed_ok} pair_fail={engine.metrics.pair_consumed_fail} '
|
||||
f'success={engine.metrics.success_count}')
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_random_with_cancellations(self):
|
||||
"""随机事件 + 随机取消 Worker,不变量仍成立。"""
|
||||
engine = PropertyEngine(seed=123)
|
||||
stop = asyncio.Event()
|
||||
errors = []
|
||||
|
||||
async def s_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.s_produce()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'S: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.005, 0.03))
|
||||
|
||||
async def p_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.p_produce()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'P: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.005, 0.03))
|
||||
|
||||
async def c_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.c_consume()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'C: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.005, 0.03))
|
||||
|
||||
async def cancel_loop():
|
||||
"""随机取消并重启 Worker。"""
|
||||
while not stop.is_set():
|
||||
await asyncio.sleep(engine.rng.uniform(0.5, 2.0))
|
||||
if tasks:
|
||||
idx = engine.rng.randint(0, len(tasks) - 1)
|
||||
t = tasks[idx]
|
||||
if not t.done():
|
||||
t.cancel()
|
||||
try:
|
||||
await t
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
# 重启
|
||||
role = idx % 3
|
||||
if role == 0:
|
||||
tasks[idx] = asyncio.create_task(s_loop())
|
||||
elif role == 1:
|
||||
tasks[idx] = asyncio.create_task(p_loop())
|
||||
else:
|
||||
tasks[idx] = asyncio.create_task(c_loop())
|
||||
|
||||
tasks = []
|
||||
for _ in range(2):
|
||||
tasks.append(asyncio.create_task(s_loop()))
|
||||
tasks.append(asyncio.create_task(p_loop()))
|
||||
tasks.append(asyncio.create_task(c_loop()))
|
||||
tasks.append(asyncio.create_task(cancel_loop()))
|
||||
|
||||
await asyncio.sleep(10)
|
||||
stop.set()
|
||||
for t in tasks:
|
||||
t.cancel()
|
||||
for t in tasks:
|
||||
try:
|
||||
await t
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
engine.check_invariants_final()
|
||||
assert not errors, f'不变量违反:\n' + '\n'.join(errors)
|
||||
print(f'\n[Property+Cancel] t_produced={engine.metrics.t_produced} '
|
||||
f'pair_ok={engine.metrics.pair_consumed_ok} '
|
||||
f'success={engine.metrics.success_count}')
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extreme_ttl_pressure(self):
|
||||
"""极端 TTL: 所有资源几乎立即过期,验证过期清理不破坏守恒。"""
|
||||
engine = PropertyEngine(
|
||||
t_slot_cap=4, q_slot_cap=4, q_pending_cap=6, phys_cap=3,
|
||||
t_max_age=0.05, q_max_age=0.05, # 50ms 过期
|
||||
seed=789,
|
||||
)
|
||||
stop = asyncio.Event()
|
||||
errors = []
|
||||
|
||||
async def s_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.s_produce()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'S: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.01, 0.05))
|
||||
|
||||
async def p_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.p_produce()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'P: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.01, 0.05))
|
||||
|
||||
async def c_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.c_consume()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'C: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.01, 0.05))
|
||||
|
||||
tasks = []
|
||||
for _ in range(2):
|
||||
tasks.append(asyncio.create_task(s_loop()))
|
||||
tasks.append(asyncio.create_task(p_loop()))
|
||||
tasks.append(asyncio.create_task(c_loop()))
|
||||
|
||||
await asyncio.sleep(5)
|
||||
stop.set()
|
||||
for t in tasks:
|
||||
t.cancel()
|
||||
for t in tasks:
|
||||
try:
|
||||
await t
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
engine.check_invariants_final()
|
||||
assert not errors, f'不变量违反:\n' + '\n'.join(errors)
|
||||
print(f'\n[TTL Pressure] t_produced={engine.metrics.t_produced} '
|
||||
f't_expired={engine.metrics.t_expired} q_expired={engine.metrics.q_expired} '
|
||||
f'pair_ok={engine.metrics.pair_consumed_ok}')
|
||||
1325
tests/test_register_runtime_unittest.py
Normal file
1325
tests/test_register_runtime_unittest.py
Normal file
File diff suppressed because it is too large
Load Diff
132
tests/test_runtime_log_analyzer.py
Normal file
132
tests/test_runtime_log_analyzer.py
Normal file
@@ -0,0 +1,132 @@
|
||||
import unittest
|
||||
|
||||
from tools.runtime_log_analyzer import (
|
||||
parse_monitor_lines,
|
||||
parse_solver_timelines,
|
||||
summarize_monitor_rows,
|
||||
summarize_solver_timelines,
|
||||
)
|
||||
|
||||
|
||||
class RuntimeLogAnalyzerTests(unittest.TestCase):
|
||||
def test_parses_csp_monitor_rows_and_recent_rates(self):
|
||||
text = "\n".join(
|
||||
[
|
||||
"[*] T:0 Q:3 phys:0 t_slot:3 q_slot:0 q_pend:0 "
|
||||
"t_prod:76 t_adm:76 t_exp:0 q_sent:87 q_ret:87 q_adm:79 q_exp:0 "
|
||||
"pair:76 ok:71 fail:0 rate:8.6/min #71",
|
||||
"[*] T:0 Q:3 phys:0 t_slot:3 q_slot:0 q_pend:0 "
|
||||
"t_prod:91 t_adm:91 t_exp:0 q_sent:101 q_ret:101 q_adm:94 q_exp:0 "
|
||||
"pair:91 ok:86 fail:0 rate:8.8/min #86",
|
||||
]
|
||||
)
|
||||
|
||||
rows = parse_monitor_lines(text)
|
||||
summary = summarize_monitor_rows(rows)
|
||||
|
||||
self.assertEqual(rows[0].kind, "csp")
|
||||
self.assertEqual(summary["last_ok"], 86)
|
||||
self.assertAlmostEqual(summary["last_cumulative_rate"], 8.8)
|
||||
self.assertAlmostEqual(summary["recent_ok_per_min"], 9.89, places=2)
|
||||
self.assertAlmostEqual(summary["recent_t_prod_per_min"], 9.89, places=2)
|
||||
self.assertEqual(summary["last_q_minus_t"], 3)
|
||||
self.assertEqual(summary["last_q_return_minus_t_prod"], 10)
|
||||
|
||||
def test_parses_state_machine_monitor_rows(self):
|
||||
text = "\n".join(
|
||||
[
|
||||
"[*] slots:7/8 act:7 cpu:83% avg:76%/85 mem:4841M "
|
||||
"T:4 Q:0 sent:0 got:0(0%) rate:0.0/min #0",
|
||||
"[*] slots:8/8 act:8 cpu:91% avg:88%/85 mem:5012M "
|
||||
"T:6 Q:5 sent:61 got:48(79%) rate:7.8/min #40",
|
||||
"[*] slots:6/8 act:6 cpu:100% avg:90%/85 mem:5351M "
|
||||
"T:1 Q:6 sent:68 got:68(100%) rate:8.8/min #59",
|
||||
]
|
||||
)
|
||||
|
||||
rows = parse_monitor_lines(text)
|
||||
summary = summarize_monitor_rows(rows)
|
||||
|
||||
self.assertEqual(rows[-1].kind, "state_machine")
|
||||
self.assertEqual(summary["last_ok"], 59)
|
||||
self.assertAlmostEqual(summary["last_cumulative_rate"], 8.8)
|
||||
self.assertEqual(summary["last_q_minus_t"], 5)
|
||||
self.assertEqual(summary["last_slots"], 6)
|
||||
|
||||
def test_parses_csp_v2_monitor_rows_with_admission_fields(self):
|
||||
text = (
|
||||
"[*] T:1 Q:2 phys:3 p_send:1 t_slot:4 q_slot:5 q_pend:6 "
|
||||
"p_batch:3.5 t_prog:1 q_inflight:4 "
|
||||
"s_phys:0.50/3.00 p_phys:0.20/1.40 c_phys:0.30/2.50 "
|
||||
"p_stage:0.50/0.80/0.40 c_stage:0.30/0.40/1.50 c_hot:4/1 "
|
||||
"solver_goto:0.10 solver_inject:0.20 solver_initial:0.50 "
|
||||
"solver_click:0.01 solver_wait:11.80 solver_reuse:0.75 solver_visible:0.00 "
|
||||
"t_prod:20 t_adm:18 t_exp:1 q_sent:24 q_ret:22 q_adm:20 q_exp:0 "
|
||||
"pair:17 ok:16 fail:1 rate:9.1/min #16"
|
||||
)
|
||||
|
||||
rows = parse_monitor_lines(text)
|
||||
summary = summarize_monitor_rows(rows)
|
||||
|
||||
self.assertEqual(len(rows), 1)
|
||||
self.assertEqual(rows[0].kind, "csp")
|
||||
self.assertEqual(summary["last_ok"], 16)
|
||||
self.assertEqual(summary["last_q_return_minus_t_prod"], 2)
|
||||
self.assertEqual(summary["last_solver_wait"], 11.8)
|
||||
self.assertEqual(summary["last_solver_reuse"], 0.75)
|
||||
self.assertEqual(summary["last_phys"], 3)
|
||||
self.assertEqual(summary["last_p_batch"], 3.5)
|
||||
self.assertEqual(summary["last_t_prog"], 1)
|
||||
self.assertEqual(summary["last_q_inflight"], 4)
|
||||
self.assertEqual(summary["last_s_phys_hold"], 3.0)
|
||||
self.assertEqual(summary["last_p_phys_wait"], 0.2)
|
||||
self.assertEqual(summary["last_c_phys_hold"], 2.5)
|
||||
self.assertEqual(summary["last_physical_hold_leader"], "s")
|
||||
self.assertEqual(summary["last_physical_wait_leader"], "s")
|
||||
self.assertEqual(summary["last_p_email_create"], 0.5)
|
||||
self.assertEqual(summary["last_p_page_prepare"], 0.8)
|
||||
self.assertEqual(summary["last_p_send"], 0.4)
|
||||
self.assertEqual(summary["last_c_page_acquire"], 0.3)
|
||||
self.assertEqual(summary["last_c_verify"], 0.4)
|
||||
self.assertEqual(summary["last_c_register"], 1.5)
|
||||
self.assertEqual(summary["last_c_hot_hits"], 4)
|
||||
self.assertEqual(summary["last_c_hot_misses"], 1)
|
||||
|
||||
def test_summarizes_solver_timeline_events(self):
|
||||
text = (
|
||||
'[solver_timeline] [{"t":0.0,"event":"page_trace_after_inject",'
|
||||
'"page_trace":{"created_at":0.0,"render_called_at":100.0,"render_returned_at":120.0,'
|
||||
'"token_written_at":null}},'
|
||||
'{"t":0.5,"event":"click_before","attempt":1,"token_len":0,'
|
||||
'"dom":{"widget":{"present":true,"visible":true},'
|
||||
'"turnstile_iframe_count":0,'
|
||||
'"element_at_center":{"tag":"DIV","is_iframe":false}}},'
|
||||
'{"t":3.0,"event":"click_after","attempt":1,"token_len":0,'
|
||||
'"click_call_ms":2500.0,'
|
||||
'"click_trace":{"mouse_move1_ms":500.0,"mouse_move2_ms":1500.0,'
|
||||
'"mouse_down_ms":300.0,"mouse_up_ms":100.0}},'
|
||||
'{"t":3.0,"event":"poll_start"},'
|
||||
'{"t":4.0,"event":"poll_done","ok":true,"token_len":730,'
|
||||
'"poll_attempts":2,"first_token_attempt":2,"poll_read_ms_avg":12.0,'
|
||||
'"poll_read_ms_max":20.0,'
|
||||
'"page_trace":{"created_at":0.0,"render_called_at":100.0,"render_returned_at":120.0,'
|
||||
'"token_written_at":3900.0}}]'
|
||||
)
|
||||
|
||||
timelines = parse_solver_timelines(text)
|
||||
summary = summarize_solver_timelines(timelines)
|
||||
|
||||
self.assertEqual(len(timelines), 1)
|
||||
self.assertEqual(summary["solver_timeline_count"], 1)
|
||||
self.assertEqual(summary["ok_count"], 1)
|
||||
self.assertEqual(summary["avg_click_call_ms"], 2500.0)
|
||||
self.assertEqual(summary["avg_click_mouse_move_ms"], 2000.0)
|
||||
self.assertEqual(summary["avg_render_to_token_ms"], 3800.0)
|
||||
self.assertEqual(summary["avg_token_write_to_poll_done_ms"], 100.0)
|
||||
self.assertEqual(summary["avg_poll_attempts"], 2.0)
|
||||
self.assertEqual(summary["center_iframe_hit_ratio"], 0.0)
|
||||
self.assertEqual(summary["turnstile_iframe_seen_ratio"], 0.0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
64
tests/test_shell_entrypoints.py
Normal file
64
tests/test_shell_entrypoints.py
Normal file
@@ -0,0 +1,64 @@
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def _entrypoint_workspace(tmp_path):
|
||||
for name in ("start.sh", "auth-service.sh", "setup.sh"):
|
||||
shutil.copy2(ROOT / name, tmp_path / name)
|
||||
scripts = tmp_path / "scripts"
|
||||
scripts.mkdir()
|
||||
shutil.copy2(ROOT / "scripts" / "ensure_runtime.sh", scripts / "ensure_runtime.sh")
|
||||
python = tmp_path / ".venv" / "bin" / "python"
|
||||
python.parent.mkdir(parents=True)
|
||||
python.write_text(
|
||||
"#!/bin/sh\nprintf '%s\\n' \"$@\" > \"$ENTRY_CAPTURE\"\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
python.chmod(0o755)
|
||||
return tmp_path
|
||||
|
||||
|
||||
def _run_entry(workspace, script, *args):
|
||||
capture = workspace / "capture.txt"
|
||||
environment = {**os.environ, "ENTRY_CAPTURE": str(capture)}
|
||||
completed = subprocess.run(
|
||||
["bash", script, *args],
|
||||
cwd=workspace,
|
||||
env=environment,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
assert completed.returncode == 0, completed.stderr
|
||||
return capture.read_text(encoding="utf-8").splitlines()
|
||||
|
||||
|
||||
def test_start_dispatches_registration_and_email_as_independent_modules(tmp_path):
|
||||
workspace = _entrypoint_workspace(tmp_path)
|
||||
(workspace / ".env").write_text("EMAIL_MODE=tempmail\n", encoding="utf-8")
|
||||
|
||||
assert _run_entry(workspace, "start.sh", "--target", "3") == [
|
||||
"-m",
|
||||
"grok_register.register",
|
||||
"--target",
|
||||
"3",
|
||||
]
|
||||
assert _run_entry(workspace, "start.sh", "--email-service", "--port", "9090") == [
|
||||
"-m",
|
||||
"grok_register.email_server",
|
||||
"--port",
|
||||
"9090",
|
||||
]
|
||||
|
||||
|
||||
def test_auth_service_keeps_the_supported_python_module_internal(tmp_path):
|
||||
workspace = _entrypoint_workspace(tmp_path)
|
||||
assert _run_entry(workspace, "auth-service.sh", "--debug") == [
|
||||
"-m",
|
||||
"xai_enroller.service",
|
||||
"--debug",
|
||||
]
|
||||
411
tests/test_stress.py
Normal file
411
tests/test_stress.py
Normal file
@@ -0,0 +1,411 @@
|
||||
"""
|
||||
Layer 4: 压力测试
|
||||
|
||||
真实 asyncio 并发跑较长时间,检查:
|
||||
- pending 泄漏
|
||||
- 库存卡死
|
||||
- claim 饥饿
|
||||
- 等待时间周期性尖峰
|
||||
- slot 守恒在长时间运行后仍成立
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
import statistics
|
||||
import pytest
|
||||
|
||||
from grok_register.core.envelope import ResourceEnvelope
|
||||
from grok_register.core.inventory import Inventory
|
||||
from grok_register.core.observer import Metrics
|
||||
from tests.fakes import Conservation, EventLog
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# S/P/C 真实并发 Worker (使用 fake 服务)
|
||||
# ═══════════════════════════════════════════
|
||||
|
||||
async def stress_s_worker(wid, t_gen, inventory, phys, t_slot, metrics, stop):
|
||||
"""压力测试 S_Worker。"""
|
||||
while not stop.is_set():
|
||||
await phys.acquire()
|
||||
try:
|
||||
token = await t_gen.generate()
|
||||
finally:
|
||||
phys.release()
|
||||
|
||||
if token is None:
|
||||
metrics.t_discarded += 1
|
||||
continue
|
||||
|
||||
metrics.t_produced += 1
|
||||
now = time.time()
|
||||
try:
|
||||
t_env = await ResourceEnvelope.create_with_slot(
|
||||
'T', token, t_slot, expires_at=now + 60
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
metrics.t_discarded += 1
|
||||
raise
|
||||
|
||||
try:
|
||||
await inventory.put_t(t_env)
|
||||
except Exception:
|
||||
t_env.discard(reason='put_fail')
|
||||
metrics.t_discarded += 1
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
|
||||
async def stress_p_worker(wid, email_svc, inventory, phys, q_pend, q_slot, metrics, stop):
|
||||
"""压力测试 P_Worker。"""
|
||||
loop = asyncio.get_event_loop()
|
||||
while not stop.is_set():
|
||||
await q_pend.acquire()
|
||||
_pend_released = False
|
||||
try:
|
||||
await phys.acquire()
|
||||
try:
|
||||
handle, email, password = await email_svc.create_email()
|
||||
sent = True
|
||||
finally:
|
||||
phys.release()
|
||||
|
||||
if not sent:
|
||||
q_pend.release()
|
||||
_pend_released = True
|
||||
continue
|
||||
|
||||
metrics.q_sent += 1
|
||||
|
||||
try:
|
||||
code = await asyncio.wait_for(email_svc.poll_code(handle), timeout=95)
|
||||
except (asyncio.TimeoutError, Exception):
|
||||
code = None
|
||||
|
||||
if code is None:
|
||||
metrics.q_discarded += 1
|
||||
q_pend.release()
|
||||
_pend_released = True
|
||||
continue
|
||||
|
||||
metrics.q_returned += 1
|
||||
now = time.time()
|
||||
q_env = await ResourceEnvelope.create_with_slot(
|
||||
'Q', {'email': email, 'password': password, 'code': code},
|
||||
q_slot, expires_at=now + 60,
|
||||
)
|
||||
|
||||
try:
|
||||
await inventory.put_q(q_env)
|
||||
except Exception:
|
||||
q_env.discard(reason='put_fail')
|
||||
metrics.q_discarded += 1
|
||||
|
||||
except asyncio.CancelledError:
|
||||
if not _pend_released:
|
||||
q_pend.release()
|
||||
_pend_released = True
|
||||
raise
|
||||
except Exception:
|
||||
metrics.q_discarded += 1
|
||||
finally:
|
||||
if not _pend_released:
|
||||
q_pend.release()
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
|
||||
async def stress_c_worker(wid, register_api, inventory, phys, metrics, stop, file_lock):
|
||||
"""压力测试 C_Worker。"""
|
||||
claim_wait_times = []
|
||||
while not stop.is_set():
|
||||
t0 = time.time()
|
||||
try:
|
||||
async with inventory.claim_pair() as pair:
|
||||
claim_wait_times.append(time.time() - t0)
|
||||
t_val = pair.t.value
|
||||
q_val = pair.q.value
|
||||
|
||||
await phys.acquire()
|
||||
try:
|
||||
sso = await register_api.register(
|
||||
q_val['email'], q_val['password'], q_val['code'], t_val
|
||||
)
|
||||
if sso:
|
||||
metrics.pair_consumed_ok += 1
|
||||
metrics.success_count += 1
|
||||
else:
|
||||
metrics.pair_consumed_fail += 1
|
||||
finally:
|
||||
phys.release()
|
||||
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
metrics.pair_consumed_fail += 1
|
||||
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
return claim_wait_times
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
class TestStress:
|
||||
"""压力测试: 长时间真实并发。"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_steady_state_30s(self):
|
||||
"""稳态运行 30 秒,检查守恒、无泄漏、无饥饿。"""
|
||||
from tests.fakes import FakeTurnstile, FakeEmailService, FakeRegisterAPI
|
||||
|
||||
t_cap, q_cap, pend_cap, phys_cap = 16, 16, 24, 8
|
||||
t_slot_sem = asyncio.Semaphore(t_cap)
|
||||
q_slot_sem = asyncio.Semaphore(q_cap)
|
||||
q_pend_sem = asyncio.Semaphore(pend_cap)
|
||||
phys_sem = asyncio.Semaphore(phys_cap)
|
||||
|
||||
ft = FakeTurnstile(delay=0.005, fail_rate=0.05)
|
||||
es = FakeEmailService()
|
||||
fra = FakeRegisterAPI(success_rate=0.9, delay=0.005)
|
||||
|
||||
metrics = Metrics()
|
||||
inv = Inventory(metrics=metrics)
|
||||
cons = Conservation(t_slot_sem, q_slot_sem, q_pend_sem, inv)
|
||||
stop = asyncio.Event()
|
||||
fl = asyncio.Lock()
|
||||
|
||||
tasks = []
|
||||
all_claim_waits = []
|
||||
|
||||
# S workers
|
||||
for i in range(4):
|
||||
tasks.append(asyncio.create_task(
|
||||
stress_s_worker(i, ft, inv, phys_sem, t_slot_sem, metrics, stop)
|
||||
))
|
||||
# P workers
|
||||
for i in range(6):
|
||||
tasks.append(asyncio.create_task(
|
||||
stress_p_worker(i, es, inv, phys_sem, q_pend_sem, q_slot_sem, metrics, stop)
|
||||
))
|
||||
# C workers
|
||||
for i in range(4):
|
||||
tasks.append(asyncio.create_task(
|
||||
stress_c_worker(i, fra, inv, phys_sem, metrics, stop, fl)
|
||||
))
|
||||
|
||||
# 监控 + 守恒检查
|
||||
async def monitor():
|
||||
snapshots = []
|
||||
while not stop.is_set():
|
||||
await asyncio.sleep(2)
|
||||
try:
|
||||
cons.check_t_conservation(t_cap)
|
||||
cons.check_q_conservation(q_cap)
|
||||
# Q_Pending: 只做软检查(非负、不超限),不做守恒检查
|
||||
# 因为 P worker 正常 acquire/release 期间 free 值会波动
|
||||
assert q_pend_sem._value >= 0, 'Q pending 为负'
|
||||
assert q_pend_sem._value <= pend_cap, 'Q pending 超限'
|
||||
except AssertionError as e:
|
||||
snapshots.append(f'INVARIANT VIOLATION: {e}')
|
||||
snapshots.append(
|
||||
f't={metrics.t_produced} q_sent={metrics.q_sent} q_ret={metrics.q_returned} '
|
||||
f'pair_ok={metrics.pair_consumed_ok} inv_t={inv.t_depth} inv_q={inv.q_depth} '
|
||||
f'phys={phys_sem._value} t_slot={t_slot_sem._value} q_slot={q_slot_sem._value}'
|
||||
)
|
||||
return snapshots
|
||||
|
||||
mon_task = asyncio.create_task(monitor())
|
||||
|
||||
await asyncio.sleep(30)
|
||||
stop.set()
|
||||
for t in tasks:
|
||||
t.cancel()
|
||||
for t in tasks:
|
||||
try:
|
||||
await t
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
snapshots = await mon_task
|
||||
|
||||
# 最终守恒检查(所有 Worker 已停止)
|
||||
cons.check_t_conservation(t_cap)
|
||||
cons.check_q_conservation(q_cap)
|
||||
# Q_Pending: 软检查
|
||||
assert q_pend_sem._value >= 0 and q_pend_sem._value <= pend_cap
|
||||
|
||||
# 检查是否有 INVARIANT VIOLATION
|
||||
violations = [s for s in snapshots if 'INVARIANT' in s]
|
||||
assert not violations, f'守恒违反:\n' + '\n'.join(violations)
|
||||
|
||||
print(f'\n[Stress 30s] t_produced={metrics.t_produced} '
|
||||
f'q_sent={metrics.q_sent} q_returned={metrics.q_returned} '
|
||||
f'pair_ok={metrics.pair_consumed_ok} pair_fail={metrics.pair_consumed_fail} '
|
||||
f'success={metrics.success_count}')
|
||||
print(f'Snapshots ({len(snapshots)}):')
|
||||
for s in snapshots[-5:]:
|
||||
print(f' {s}')
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_claim_starvation(self):
|
||||
"""检查 C 不会永远等不到 pair (claim 饥饿)。"""
|
||||
from tests.fakes import FakeTurnstile, FakeEmailService, FakeRegisterAPI
|
||||
|
||||
t_cap, q_cap, pend_cap, phys_cap = 4, 4, 6, 2
|
||||
t_slot_sem = asyncio.Semaphore(t_cap)
|
||||
q_slot_sem = asyncio.Semaphore(q_cap)
|
||||
q_pend_sem = asyncio.Semaphore(pend_cap)
|
||||
phys_sem = asyncio.Semaphore(phys_cap)
|
||||
|
||||
ft = FakeTurnstile(delay=0.002, fail_rate=0.0)
|
||||
es = FakeEmailService()
|
||||
fra = FakeRegisterAPI(success_rate=1.0, delay=0.002)
|
||||
|
||||
metrics = Metrics()
|
||||
inv = Inventory(metrics=metrics)
|
||||
stop = asyncio.Event()
|
||||
fl = asyncio.Lock()
|
||||
|
||||
claim_wait_samples = []
|
||||
last_success_time = time.time()
|
||||
|
||||
async def c_worker(wid):
|
||||
nonlocal last_success_time
|
||||
while not stop.is_set():
|
||||
t0 = time.time()
|
||||
try:
|
||||
async with inv.claim_pair() as pair:
|
||||
wait = time.time() - t0
|
||||
claim_wait_samples.append(wait)
|
||||
|
||||
await phys_sem.acquire()
|
||||
try:
|
||||
sso = await fra.register(
|
||||
pair.q.value['email'], pair.q.value['password'],
|
||||
pair.q.value['code'], pair.t.value
|
||||
)
|
||||
if sso:
|
||||
metrics.pair_consumed_ok += 1
|
||||
metrics.success_count += 1
|
||||
last_success_time = time.time()
|
||||
finally:
|
||||
phys_sem.release()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception:
|
||||
metrics.pair_consumed_fail += 1
|
||||
await asyncio.sleep(0.005)
|
||||
|
||||
tasks = []
|
||||
for i in range(3):
|
||||
tasks.append(asyncio.create_task(stress_s_worker(i, ft, inv, phys_sem, t_slot_sem, metrics, stop)))
|
||||
for i in range(4):
|
||||
tasks.append(asyncio.create_task(stress_p_worker(i, es, inv, phys_sem, q_pend_sem, q_slot_sem, metrics, stop)))
|
||||
for i in range(3):
|
||||
tasks.append(asyncio.create_task(c_worker(i)))
|
||||
|
||||
await asyncio.sleep(15)
|
||||
stop.set()
|
||||
for t in tasks:
|
||||
t.cancel()
|
||||
for t in tasks:
|
||||
try:
|
||||
await t
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# 检查: claim 等待时间不应有极端尖峰
|
||||
if claim_wait_samples:
|
||||
p50 = statistics.median(claim_wait_samples)
|
||||
p95 = sorted(claim_wait_samples)[int(len(claim_wait_samples) * 0.95)]
|
||||
p99 = sorted(claim_wait_samples)[int(len(claim_wait_samples) * 0.99)]
|
||||
print(f'\n[Starvation] claim_wait p50={p50:.3f}s p95={p95:.3f}s p99={p99:.3f}s')
|
||||
print(f' samples={len(claim_wait_samples)} pair_ok={metrics.pair_consumed_ok}')
|
||||
# p99 不应超过 5 秒(在 fake 服务下)
|
||||
assert p99 < 5.0, f'claim 饥饿: p99={p99:.3f}s 太高'
|
||||
|
||||
# 检查: 最近成功消费不应太久
|
||||
time_since_last = time.time() - last_success_time
|
||||
assert time_since_last < 5.0, f'太久没有成功消费: {time_since_last:.1f}s'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_high_cancellation_churn(self):
|
||||
"""高取消翻转: Worker 频繁取消重启,系统仍保持守恒。"""
|
||||
from tests.fakes import FakeTurnstile, FakeEmailService, FakeRegisterAPI
|
||||
|
||||
t_cap, q_cap, pend_cap, phys_cap = 8, 8, 12, 4
|
||||
t_slot_sem = asyncio.Semaphore(t_cap)
|
||||
q_slot_sem = asyncio.Semaphore(q_cap)
|
||||
q_pend_sem = asyncio.Semaphore(pend_cap)
|
||||
phys_sem = asyncio.Semaphore(phys_cap)
|
||||
|
||||
ft = FakeTurnstile(delay=0.003, fail_rate=0.1)
|
||||
es = FakeEmailService()
|
||||
fra = FakeRegisterAPI(success_rate=0.85, delay=0.003)
|
||||
|
||||
metrics = Metrics()
|
||||
inv = Inventory(metrics=metrics)
|
||||
cons = Conservation(t_slot_sem, q_slot_sem, q_pend_sem, inv)
|
||||
stop = asyncio.Event()
|
||||
fl = asyncio.Lock()
|
||||
|
||||
tasks = []
|
||||
violations = []
|
||||
|
||||
def spawn_worker(kind, idx):
|
||||
if kind == 'S':
|
||||
return asyncio.create_task(stress_s_worker(idx, ft, inv, phys_sem, t_slot_sem, metrics, stop))
|
||||
elif kind == 'P':
|
||||
return asyncio.create_task(stress_p_worker(idx, es, inv, phys_sem, q_pend_sem, q_slot_sem, metrics, stop))
|
||||
else:
|
||||
return asyncio.create_task(stress_c_worker(idx, fra, inv, phys_sem, metrics, stop, fl))
|
||||
|
||||
# 初始 workers
|
||||
for i in range(3):
|
||||
tasks.append(spawn_worker('S', i))
|
||||
tasks.append(spawn_worker('P', i))
|
||||
tasks.append(spawn_worker('C', i))
|
||||
|
||||
async def churn():
|
||||
import random
|
||||
while not stop.is_set():
|
||||
await asyncio.sleep(0.3)
|
||||
# 随机取消一个,重启一个
|
||||
if tasks:
|
||||
idx = random.randint(0, len(tasks) - 1)
|
||||
t = tasks[idx]
|
||||
if not t.done():
|
||||
t.cancel()
|
||||
tasks[idx] = spawn_worker(
|
||||
random.choice(['S', 'P', 'C']),
|
||||
random.randint(0, 9)
|
||||
)
|
||||
# 定期守恒检查(软检查)
|
||||
try:
|
||||
cons.check_t_conservation(t_cap)
|
||||
cons.check_q_conservation(q_cap)
|
||||
assert q_pend_sem._value >= 0 and q_pend_sem._value <= pend_cap
|
||||
except AssertionError as e:
|
||||
violations.append(str(e))
|
||||
|
||||
churn_task = asyncio.create_task(churn())
|
||||
tasks.append(churn_task)
|
||||
|
||||
await asyncio.sleep(15)
|
||||
stop.set()
|
||||
for t in tasks:
|
||||
t.cancel()
|
||||
for t in tasks:
|
||||
try:
|
||||
await t
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# 最终守恒检查
|
||||
try:
|
||||
cons.check_t_conservation(t_cap)
|
||||
cons.check_q_conservation(q_cap)
|
||||
assert q_pend_sem._value >= 0 and q_pend_sem._value <= pend_cap
|
||||
except AssertionError as e:
|
||||
violations.append(f'FINAL: {e}')
|
||||
|
||||
assert not violations, f'守恒违反:\n' + '\n'.join(violations[:5])
|
||||
print(f'\n[High Churn] t_produced={metrics.t_produced} '
|
||||
f'pair_ok={metrics.pair_consumed_ok} success={metrics.success_count}')
|
||||
298
tests/test_xai_auth_pipeline.py
Normal file
298
tests/test_xai_auth_pipeline.py
Normal file
@@ -0,0 +1,298 @@
|
||||
import asyncio
|
||||
import threading
|
||||
|
||||
from xai_enroller.auth_pipeline import AuthPipeline
|
||||
from xai_enroller.ledger import Ledger
|
||||
from xai_enroller.models import (
|
||||
AuthorizationResult,
|
||||
AuthorizationStatus,
|
||||
DeviceFlow,
|
||||
JobStatus,
|
||||
OAuthCredential,
|
||||
PipelineState,
|
||||
SinkReceipt,
|
||||
SourceRecord,
|
||||
)
|
||||
|
||||
|
||||
class EmptySource:
|
||||
async def records(self):
|
||||
if False:
|
||||
yield None
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
|
||||
class NoopExecutor:
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
|
||||
def _pipeline(tmp_path):
|
||||
return AuthPipeline(
|
||||
source=EmptySource(),
|
||||
protocol=object(),
|
||||
executor=NoopExecutor(),
|
||||
sink=object(),
|
||||
ledger=Ledger(tmp_path / "ledger.db", b"salt"),
|
||||
)
|
||||
|
||||
|
||||
def test_cancel_active_is_idempotent_while_cleanup_is_in_progress(tmp_path):
|
||||
async def scenario():
|
||||
pipeline = _pipeline(tmp_path)
|
||||
|
||||
async def wait_forever():
|
||||
await asyncio.Future()
|
||||
|
||||
task = asyncio.create_task(wait_forever())
|
||||
pipeline._authorization_task = task
|
||||
pipeline._authorization_cancellable = True
|
||||
assert await pipeline.cancel_active()
|
||||
assert not await pipeline.cancel_active()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
class SingleSource:
|
||||
async def records(self):
|
||||
yield SourceRecord("stock", "opaque")
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
|
||||
class ImmediateProtocol:
|
||||
async def start_device_flow(self):
|
||||
return DeviceFlow(
|
||||
"device",
|
||||
"code",
|
||||
"https://accounts.x.ai/oauth2/device",
|
||||
600,
|
||||
0,
|
||||
)
|
||||
|
||||
async def poll_token(self, **_kwargs):
|
||||
return OAuthCredential(
|
||||
"access",
|
||||
"refresh",
|
||||
None,
|
||||
"Bearer",
|
||||
3600,
|
||||
"later",
|
||||
"now",
|
||||
"subject",
|
||||
"https://auth.x.ai/oauth2/token",
|
||||
)
|
||||
|
||||
|
||||
class ImmediateExecutor:
|
||||
def __init__(self):
|
||||
self.calls = 0
|
||||
|
||||
async def start(self):
|
||||
return None
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
async def confirm(self, *_args):
|
||||
self.calls += 1
|
||||
return AuthorizationResult(AuthorizationStatus.AUTHORIZED, "confirmed")
|
||||
|
||||
|
||||
class BlockingThreadSink:
|
||||
def __init__(self):
|
||||
self.entered = threading.Event()
|
||||
self.release = threading.Event()
|
||||
self.pipeline = None
|
||||
self.calls = 0
|
||||
|
||||
def _store(self):
|
||||
self.calls += 1
|
||||
self.entered.set()
|
||||
self.release.wait(2)
|
||||
return SinkReceipt("opaque")
|
||||
|
||||
async def store(self, _credential):
|
||||
receipt = await asyncio.to_thread(self._store)
|
||||
self.pipeline.request_stop()
|
||||
return receipt
|
||||
|
||||
|
||||
def test_sink_and_ledger_commit_cannot_be_cancelled_mid_commit(tmp_path):
|
||||
async def scenario():
|
||||
executor = ImmediateExecutor()
|
||||
sink = BlockingThreadSink()
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
events = []
|
||||
pipeline = AuthPipeline(
|
||||
source=SingleSource(),
|
||||
protocol=ImmediateProtocol(),
|
||||
executor=executor,
|
||||
sink=sink,
|
||||
ledger=ledger,
|
||||
timeout=2,
|
||||
event_callback=lambda kind, data: events.append((kind, data)),
|
||||
)
|
||||
sink.pipeline = pipeline
|
||||
run = asyncio.create_task(pipeline.run())
|
||||
assert await asyncio.to_thread(sink.entered.wait, 1)
|
||||
assert not await pipeline.cancel_active()
|
||||
sink.release.set()
|
||||
await asyncio.wait_for(run, 2)
|
||||
assert executor.calls == 1
|
||||
assert sink.calls == 1
|
||||
assert ledger.aggregate_counts()["imported_unique"] == 1
|
||||
started = next(data for kind, data in events if kind == "authorization_started")
|
||||
imported = next(
|
||||
data
|
||||
for kind, data in events
|
||||
if kind == "result" and data["status"] == "imported"
|
||||
)
|
||||
assert imported["task_number"] == started["task_number"]
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_pending_total_uses_current_snapshot_set_difference(tmp_path):
|
||||
class SnapshotSource(EmptySource):
|
||||
snapshot_fingerprints = frozenset({"a", "b", "c"})
|
||||
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
imported = ledger.start_fingerprint("b")
|
||||
ledger.finish(imported, JobStatus.IMPORTED, "imported", "receipt-b")
|
||||
outside = ledger.start_fingerprint("outside")
|
||||
ledger.finish(outside, JobStatus.IMPORTED, "imported", "receipt-outside")
|
||||
pipeline = AuthPipeline(
|
||||
source=SnapshotSource(),
|
||||
protocol=object(),
|
||||
executor=NoopExecutor(),
|
||||
sink=object(),
|
||||
ledger=ledger,
|
||||
)
|
||||
|
||||
assert pipeline.status()["pending_total"] == 2
|
||||
assert pipeline.status()["five_minute_imports_per_minute"] is None
|
||||
assert pipeline.status()["lifetime_imports_per_minute"] is None
|
||||
|
||||
|
||||
def test_pending_total_is_unknown_before_first_valid_snapshot(tmp_path):
|
||||
pipeline = _pipeline(tmp_path)
|
||||
|
||||
assert pipeline.status()["pending_total"] is None
|
||||
|
||||
|
||||
def test_rate_limit_gate_grows_cooldown_on_consecutive_trips():
|
||||
class Clock:
|
||||
def __init__(self):
|
||||
self.now = 0.0
|
||||
|
||||
def __call__(self):
|
||||
return self.now
|
||||
|
||||
async def scenario():
|
||||
from xai_enroller.auth_pipeline import GlobalRateLimitGate
|
||||
|
||||
clock = Clock()
|
||||
gate = GlobalRateLimitGate(clock=clock, base_cooldown=60.0)
|
||||
assert await gate.rate_limited()
|
||||
assert gate.COOLDOWN_SECONDS == 60.0
|
||||
assert gate.snapshot()["cooldown_seconds"] == 60.0
|
||||
|
||||
resume = asyncio.Event()
|
||||
resume.set()
|
||||
clock.now = 60.0
|
||||
probe = await gate.wait_for_permission(resume)
|
||||
assert probe is not None
|
||||
assert await gate.rate_limited(probe)
|
||||
assert gate.COOLDOWN_SECONDS == 90.0
|
||||
assert gate.snapshot()["cooldown_remaining_seconds"] == 90.0
|
||||
|
||||
clock.now = 150.0
|
||||
probe = await gate.wait_for_permission(resume)
|
||||
elapsed = await gate.authorized(probe)
|
||||
assert elapsed == 150.0
|
||||
assert gate.COOLDOWN_SECONDS == 60.0
|
||||
assert gate.snapshot()["rate_limit_trips"] == 0
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_minimum_start_interval_raises_after_rate_limit_and_decays():
|
||||
from xai_enroller.auth_pipeline import MinimumStartInterval
|
||||
|
||||
def approx(value, rel=1e-6):
|
||||
class Approx:
|
||||
def __eq__(self, other):
|
||||
return abs(float(other) - float(value)) <= rel * max(1.0, abs(value))
|
||||
|
||||
def __repr__(self):
|
||||
return f"approx({value})"
|
||||
|
||||
return Approx()
|
||||
|
||||
interval = MinimumStartInterval(10.0)
|
||||
assert interval.seconds == 10.0
|
||||
raised = interval.note_rate_limited()
|
||||
assert raised == 18.0
|
||||
raised = interval.note_rate_limited()
|
||||
assert raised == approx(18.0 * 1.55)
|
||||
for _ in range(MinimumStartInterval.SUCCESS_STREAK_TO_DECAY):
|
||||
interval.note_success()
|
||||
assert interval.seconds < raised
|
||||
# Keep decaying until the base floor is restored.
|
||||
for _ in range(40):
|
||||
interval.note_success()
|
||||
assert interval.seconds == 10.0
|
||||
|
||||
|
||||
def test_queued_source_imported_while_paused_is_not_authorized(tmp_path):
|
||||
async def scenario():
|
||||
executor = ImmediateExecutor()
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
pipeline = AuthPipeline(
|
||||
source=EmptySource(),
|
||||
protocol=ImmediateProtocol(),
|
||||
executor=executor,
|
||||
sink=object(),
|
||||
ledger=ledger,
|
||||
timeout=2,
|
||||
)
|
||||
source = SourceRecord("stock", "opaque")
|
||||
fingerprint = ledger.fingerprint(source.source_id)
|
||||
|
||||
assert await pipeline.admit(source)
|
||||
pipeline.pause()
|
||||
await pipeline.rate_gate.rate_limited()
|
||||
run = asyncio.create_task(pipeline.run())
|
||||
try:
|
||||
async with asyncio.timeout(1):
|
||||
while pipeline._authorization_task is None:
|
||||
await asyncio.sleep(0)
|
||||
|
||||
external_job_id = ledger.start_fingerprint(fingerprint)
|
||||
ledger.finish(
|
||||
external_job_id,
|
||||
JobStatus.IMPORTED,
|
||||
"imported",
|
||||
"external-receipt",
|
||||
)
|
||||
pipeline.resume()
|
||||
|
||||
async with asyncio.timeout(1):
|
||||
while (
|
||||
executor.calls == 0
|
||||
and pipeline._states.get(fingerprint) is not PipelineState.IMPORTED
|
||||
):
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert executor.calls == 0
|
||||
assert pipeline._states[fingerprint] is PipelineState.IMPORTED
|
||||
finally:
|
||||
pipeline.request_stop()
|
||||
await asyncio.wait_for(run, 2)
|
||||
|
||||
asyncio.run(scenario())
|
||||
391
tests/test_xai_auth_service.py
Normal file
391
tests/test_xai_auth_service.py
Normal file
@@ -0,0 +1,391 @@
|
||||
import asyncio
|
||||
import io
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from xai_enroller.models import JobStatus
|
||||
from xai_enroller.service import (
|
||||
AuthServiceSettings,
|
||||
AuthPipelineRunner,
|
||||
EventTerminal,
|
||||
InteractiveCommandPrompt,
|
||||
resolve_auth_log_mode,
|
||||
)
|
||||
from xai_enroller.ledger import Ledger
|
||||
from xai_enroller.inventory import InventoryError
|
||||
|
||||
|
||||
def test_auth_log_mode_prefers_cli_and_rejects_invalid_values():
|
||||
assert resolve_auth_log_mode([], {}) == "user"
|
||||
assert resolve_auth_log_mode(
|
||||
[], {"XAI_AUTH_SERVICE_LOG_MODE": "debug"}
|
||||
) == "debug"
|
||||
assert resolve_auth_log_mode(
|
||||
["--debug"], {"XAI_AUTH_SERVICE_LOG_MODE": "user"}
|
||||
) == "debug"
|
||||
with pytest.raises(ValueError, match="XAI_AUTH_SERVICE_LOG_MODE"):
|
||||
resolve_auth_log_mode([], {"XAI_AUTH_SERVICE_LOG_MODE": "verbose"})
|
||||
|
||||
|
||||
def test_user_terminal_reports_progress_without_internal_or_secret_fields():
|
||||
messages = []
|
||||
terminal = EventTerminal(mode="user", output=messages.append)
|
||||
terminal.emit(
|
||||
(
|
||||
"startup",
|
||||
{
|
||||
"available": 7,
|
||||
"claimed": 3,
|
||||
"destination": "authenticated/",
|
||||
"source_kind": "local",
|
||||
"ssh_host": "do-not-print.example",
|
||||
},
|
||||
)
|
||||
)
|
||||
terminal.emit(
|
||||
(
|
||||
"authorization_started",
|
||||
{
|
||||
"task_number": 3,
|
||||
"attempt_number": 1,
|
||||
"pending_total": 41,
|
||||
"source_queue": 64,
|
||||
"source_id": "secret@example.test",
|
||||
},
|
||||
)
|
||||
)
|
||||
terminal.emit(
|
||||
(
|
||||
"result",
|
||||
{
|
||||
"status": "imported",
|
||||
"reason": "imported",
|
||||
"task_number": 3,
|
||||
"five_minute_imports_per_minute": 4.25,
|
||||
"lifetime_imports_per_minute": 1.75,
|
||||
"imported_unique": 120,
|
||||
"available": 117,
|
||||
"source_id": "secret@example.test",
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
assert messages == [
|
||||
"[✓] 本地认证服务已启动 | 来源 本机 | 输出 authenticated/ | 待处理 — | 可用 7",
|
||||
"[→] 开始认证 #3 | 待处理 41",
|
||||
"[✓] 认证成功 #3 | 运行平均 1.75/分 | 累计 120 | 可用 117",
|
||||
]
|
||||
rendered = "\n".join(messages)
|
||||
assert "source_queue" not in rendered
|
||||
assert "secret@example.test" not in rendered
|
||||
assert "do-not-print.example" not in rendered
|
||||
|
||||
|
||||
def test_user_status_distinguishes_unknown_runtime_rate_from_zero_rate():
|
||||
base = {
|
||||
"state": "running",
|
||||
"pending_total": None,
|
||||
"active_stage": "idle",
|
||||
"imported_unique": 0,
|
||||
"available": 0,
|
||||
"claimed": 0,
|
||||
"cooldown": False,
|
||||
}
|
||||
messages = []
|
||||
terminal = EventTerminal(mode="user", output=messages.append)
|
||||
terminal.emit(("status", {**base, "lifetime_imports_per_minute": None}))
|
||||
terminal.emit(("status", {**base, "lifetime_imports_per_minute": 0.0}))
|
||||
|
||||
assert "运行平均 —" in messages[0]
|
||||
assert "运行平均 0.00/分" in messages[1]
|
||||
assert "source_queue" not in messages[0]
|
||||
|
||||
|
||||
def test_user_terminal_omits_missing_task_number_before_authorization_starts():
|
||||
messages = []
|
||||
terminal = EventTerminal(mode="user", output=messages.append)
|
||||
|
||||
terminal.emit(
|
||||
(
|
||||
"result",
|
||||
{
|
||||
"status": "failed",
|
||||
"reason": "device_flow_failed",
|
||||
"task_number": None,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
assert messages == ["[✗] 认证未完成 | 暂时失败，将自动重试"]
|
||||
assert "None" not in messages[0]
|
||||
|
||||
|
||||
def test_debug_terminal_keeps_aggregate_diagnostics_and_sanitizes_unknown_events():
|
||||
messages = []
|
||||
terminal = EventTerminal(mode="debug", output=messages.append)
|
||||
terminal.emit(
|
||||
(
|
||||
"status",
|
||||
{
|
||||
"state": "running",
|
||||
"source_queue": 2,
|
||||
"prepared_queue": 1,
|
||||
"completion_queue": 0,
|
||||
"active_stage": "authorization",
|
||||
"retry_waiting": 1,
|
||||
"next_retry_seconds": 5.0,
|
||||
"authorization_starts": 3,
|
||||
"cooldown": False,
|
||||
"cooldown_remaining_seconds": 0.0,
|
||||
"probe_in_flight": False,
|
||||
"min_authorization_interval_seconds": 10.0,
|
||||
"pacing_remaining_seconds": 2.0,
|
||||
"imported_unique": 2,
|
||||
"attempted_unique": 3,
|
||||
"rate_limited": 1,
|
||||
"five_minute_imports_per_minute": 2.0,
|
||||
"lifetime_imports_per_minute": 1.0,
|
||||
"available": 2,
|
||||
"claiming": 0,
|
||||
"claimed": 0,
|
||||
},
|
||||
)
|
||||
)
|
||||
terminal.emit(
|
||||
(
|
||||
"future_event",
|
||||
{"reason": "internal_error", "token": "do-not-print-token"},
|
||||
)
|
||||
)
|
||||
|
||||
assert "queues=2/1/0" in messages[0]
|
||||
assert messages[1] == "• debug event=future_event reason=internal_error"
|
||||
assert "do-not-print-token" not in "\n".join(messages)
|
||||
|
||||
|
||||
def test_terminal_output_failure_does_not_escape():
|
||||
def broken_output(_message):
|
||||
raise OSError("closed terminal")
|
||||
|
||||
terminal = EventTerminal(mode="user", output=broken_output)
|
||||
terminal.emit(("service_stopped", {}))
|
||||
|
||||
|
||||
def test_interactive_prompt_restores_partially_typed_take_command_after_events():
|
||||
output = io.StringIO()
|
||||
commands = []
|
||||
prompt = InteractiveCommandPrompt(
|
||||
output=output,
|
||||
interactive=True,
|
||||
prompt="认证> ",
|
||||
)
|
||||
prompt.start(commands.append)
|
||||
prompt.feed("take 12")
|
||||
|
||||
prompt.write_event("[↻] 发现新账号 3")
|
||||
|
||||
assert output.getvalue().endswith("[↻] 发现新账号 3\n认证> take 12")
|
||||
|
||||
prompt.feed("\r")
|
||||
|
||||
assert commands == ["take 12"]
|
||||
assert output.getvalue().endswith("\n认证> ")
|
||||
|
||||
|
||||
def test_auth_service_configuration_errors_are_actionable_without_traceback(tmp_path):
|
||||
environment = os.environ.copy()
|
||||
environment.pop("XAI_AUTH_SERVICE_SSH_HOST", None)
|
||||
environment.pop("XAI_AUTH_SERVICE_LOG_MODE", None)
|
||||
environment["XAI_AUTH_SERVICE_SOURCE"] = "ssh"
|
||||
missing = subprocess.run(
|
||||
[sys.executable, "-m", "xai_enroller.service"],
|
||||
env=environment,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
assert missing.returncode == 2
|
||||
assert "XAI_AUTH_SERVICE_SSH_HOST" in missing.stderr
|
||||
assert "docs/guides/auth-service.md" in missing.stderr
|
||||
assert "Traceback" not in missing.stderr
|
||||
|
||||
environment["XAI_AUTH_SERVICE_LOG_MODE"] = "verbose"
|
||||
invalid = subprocess.run(
|
||||
[sys.executable, "-m", "xai_enroller.service"],
|
||||
env=environment,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
assert invalid.returncode == 2
|
||||
assert "XAI_AUTH_SERVICE_LOG_MODE" in invalid.stderr
|
||||
assert "Traceback" not in invalid.stderr
|
||||
|
||||
environment["XAI_AUTH_SERVICE_LOG_MODE"] = "user"
|
||||
environment.pop("XAI_AUTH_SERVICE_SOURCE", None)
|
||||
environment["XAI_AUTH_SERVICE_SSH_HOST"] = "user@example.test"
|
||||
environment["XAI_ENROLLER_LOCAL_AUTH_DIR"] = str(tmp_path / "auth")
|
||||
environment["XAI_ENROLLER_TIMEOUT_SEC"] = "not-a-number"
|
||||
invalid_enroller_setting = subprocess.run(
|
||||
[sys.executable, "-m", "xai_enroller.service"],
|
||||
env=environment,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
assert invalid_enroller_setting.returncode == 2
|
||||
assert "XAI_ENROLLER_TIMEOUT_SEC" in invalid_enroller_setting.stderr
|
||||
assert "docs/guides/auth-service.md" in invalid_enroller_setting.stderr
|
||||
assert "Traceback" not in invalid_enroller_setting.stderr
|
||||
|
||||
|
||||
def test_auth_service_startup_failure_is_sanitized_without_traceback(tmp_path):
|
||||
blocked_destination = tmp_path / "private-output-path"
|
||||
blocked_destination.write_text("not a directory", encoding="utf-8")
|
||||
environment = os.environ.copy()
|
||||
environment["XAI_AUTH_SERVICE_SSH_HOST"] = "user@example.test"
|
||||
environment["XAI_ENROLLER_LOCAL_AUTH_DIR"] = str(blocked_destination)
|
||||
|
||||
failed = subprocess.run(
|
||||
[sys.executable, "-m", "xai_enroller.service"],
|
||||
env=environment,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
assert failed.returncode == 1
|
||||
assert "认证服务异常终止" in failed.stderr
|
||||
assert "bash auth-service.sh --debug" in failed.stderr
|
||||
assert "Traceback" not in failed.stderr
|
||||
assert str(blocked_destination) not in failed.stderr
|
||||
|
||||
|
||||
def test_ledger_recognizes_previously_imported_sources(tmp_path):
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"test-salt")
|
||||
imported = ledger.start("done@example.test")
|
||||
ledger.finish(imported, JobStatus.IMPORTED, "imported")
|
||||
failed = ledger.start("retry@example.test")
|
||||
ledger.finish(failed, JobStatus.SINK_FAILED, "sink_failed")
|
||||
|
||||
assert ledger.has_imported("done@example.test") is True
|
||||
assert ledger.has_imported("retry@example.test") is False
|
||||
|
||||
|
||||
def test_auth_service_settings_defaults_local_and_preserves_ssh_auto_detection(tmp_path):
|
||||
local = AuthServiceSettings.from_environ(
|
||||
{"XAI_AUTH_SERVICE_REGISTER_ROOT": str(tmp_path)}
|
||||
)
|
||||
remote = AuthServiceSettings.from_environ(
|
||||
{
|
||||
"XAI_AUTH_SERVICE_SSH_HOST": "ubuntu@example.test",
|
||||
"XAI_AUTH_SERVICE_SYNC_SEC": "45",
|
||||
}
|
||||
)
|
||||
|
||||
assert local.source_kind == "local"
|
||||
assert local.ssh_host is None
|
||||
assert local.register_root == str(tmp_path)
|
||||
assert remote.source_kind == "ssh"
|
||||
assert remote.ssh_host == "ubuntu@example.test"
|
||||
assert remote.sync_seconds == 45
|
||||
assert remote.remote_root == "/opt/grok-free-register"
|
||||
|
||||
|
||||
def test_auth_service_settings_explicit_source_overrides_auto_detection():
|
||||
local = AuthServiceSettings.from_environ(
|
||||
{
|
||||
"XAI_AUTH_SERVICE_SOURCE": "local",
|
||||
"XAI_AUTH_SERVICE_SSH_HOST": "ignored@example.test",
|
||||
}
|
||||
)
|
||||
assert local.source_kind == "local"
|
||||
|
||||
with pytest.raises(ValueError, match="XAI_AUTH_SERVICE_SSH_HOST"):
|
||||
AuthServiceSettings.from_environ({"XAI_AUTH_SERVICE_SOURCE": "ssh"})
|
||||
with pytest.raises(ValueError, match="XAI_AUTH_SERVICE_SOURCE"):
|
||||
AuthServiceSettings.from_environ({"XAI_AUTH_SERVICE_SOURCE": "unknown"})
|
||||
|
||||
|
||||
def test_pipeline_runner_takes_a_credential_batch_and_reports_inventory():
|
||||
class InventoryLedger:
|
||||
def inventory_counts(self):
|
||||
return {"available": 7, "claiming": 0, "claimed": 3}
|
||||
|
||||
class Pipeline:
|
||||
ledger = InventoryLedger()
|
||||
|
||||
def status(self):
|
||||
return {"state": "running"}
|
||||
|
||||
class Inventory:
|
||||
def take(self, count):
|
||||
assert count == 3
|
||||
return SimpleNamespace(
|
||||
batch_id="batch-1",
|
||||
directory=Path("/tmp/claimed/batch-1"),
|
||||
moved=3,
|
||||
note="",
|
||||
)
|
||||
|
||||
async def scenario():
|
||||
events = []
|
||||
runner = AuthPipelineRunner(Pipeline(), events.append, inventory=Inventory())
|
||||
assert await runner.handle_command("take 3") is True
|
||||
assert await runner.handle_command("s") is True
|
||||
return events
|
||||
|
||||
assert asyncio.run(scenario()) == [
|
||||
(
|
||||
"inventory_taken",
|
||||
{
|
||||
"batch_id": "batch-1",
|
||||
"directory": "/tmp/claimed/batch-1",
|
||||
"moved": 3,
|
||||
"available": 7,
|
||||
"claiming": 0,
|
||||
"claimed": 3,
|
||||
},
|
||||
),
|
||||
(
|
||||
"status",
|
||||
{
|
||||
"state": "running",
|
||||
"available": 7,
|
||||
"claiming": 0,
|
||||
"claimed": 3,
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_pipeline_runner_reports_inventory_failure_without_stopping():
|
||||
class Ledger:
|
||||
def inventory_counts(self):
|
||||
return {"available": 1, "claiming": 1, "claimed": 0}
|
||||
|
||||
class Pipeline:
|
||||
ledger = Ledger()
|
||||
|
||||
class Inventory:
|
||||
def take(self, _count):
|
||||
raise InventoryError("credential file is missing")
|
||||
|
||||
async def scenario():
|
||||
events = []
|
||||
runner = AuthPipelineRunner(Pipeline(), events.append, inventory=Inventory())
|
||||
assert await runner.handle_command("take 1") is True
|
||||
return events
|
||||
|
||||
assert asyncio.run(scenario()) == [
|
||||
(
|
||||
"inventory_error",
|
||||
{
|
||||
"reason": "credential file is missing",
|
||||
"available": 1,
|
||||
"claiming": 1,
|
||||
"claimed": 0,
|
||||
},
|
||||
)
|
||||
]
|
||||
58
tests/test_xai_credential_inventory.py
Normal file
58
tests/test_xai_credential_inventory.py
Normal file
@@ -0,0 +1,58 @@
|
||||
from xai_enroller.inventory import CredentialInventory
|
||||
from xai_enroller.ledger import Ledger
|
||||
from xai_enroller.models import JobStatus
|
||||
|
||||
|
||||
def import_credential(ledger, source, receipt):
|
||||
job_id = ledger.start(source)
|
||||
ledger.finish(job_id, JobStatus.IMPORTED, "imported", receipt)
|
||||
return job_id
|
||||
|
||||
|
||||
def test_take_moves_available_credentials_without_changing_imported_jobs(tmp_path):
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
available = tmp_path / "authenticated"
|
||||
claimed = tmp_path / "claimed"
|
||||
available.mkdir()
|
||||
first_job = import_credential(ledger, "first", "xai-first")
|
||||
second_job = import_credential(ledger, "second", "xai-second")
|
||||
(available / "xai-first.json").write_text("{}\n", encoding="utf-8")
|
||||
(available / "xai-second.json").write_text("{}\n", encoding="utf-8")
|
||||
|
||||
batch = CredentialInventory(ledger, available, claimed).take(2)
|
||||
|
||||
assert batch.moved == 2
|
||||
assert batch.note == ""
|
||||
assert not list(available.glob("*.json"))
|
||||
assert sorted(path.name for path in batch.directory.glob("*.json")) == [
|
||||
"xai-first.json",
|
||||
"xai-second.json",
|
||||
]
|
||||
assert ledger.inventory_counts() == {
|
||||
"available": 0,
|
||||
"claiming": 0,
|
||||
"claimed": 2,
|
||||
}
|
||||
assert ledger.get(first_job)["status"] == JobStatus.IMPORTED.value
|
||||
assert ledger.get(second_job)["status"] == JobStatus.IMPORTED.value
|
||||
|
||||
|
||||
def test_recover_finishes_a_batch_interrupted_before_file_move(tmp_path):
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
available = tmp_path / "authenticated"
|
||||
claimed = tmp_path / "claimed"
|
||||
available.mkdir()
|
||||
import_credential(ledger, "source", "xai-recover")
|
||||
(available / "xai-recover.json").write_text("{}\n", encoding="utf-8")
|
||||
assert ledger.claim_available(1, "batch-recover") == ["xai-recover"]
|
||||
|
||||
recovered = CredentialInventory(ledger, available, claimed).recover()
|
||||
|
||||
assert recovered == 1
|
||||
assert not (available / "xai-recover.json").exists()
|
||||
assert (claimed / "batch-recover" / "xai-recover.json").exists()
|
||||
assert ledger.inventory_counts() == {
|
||||
"available": 0,
|
||||
"claiming": 0,
|
||||
"claimed": 1,
|
||||
}
|
||||
125
tests/test_xai_enroller_config.py
Normal file
125
tests/test_xai_enroller_config.py
Normal file
@@ -0,0 +1,125 @@
|
||||
import os
|
||||
import sqlite3
|
||||
import stat
|
||||
|
||||
import pytest
|
||||
|
||||
from xai_enroller.config import Settings
|
||||
from xai_enroller.sources import FileSourceAdapter, SQLiteSourceAdapter
|
||||
|
||||
|
||||
def base_env(tmp_path, **overrides):
|
||||
values = {
|
||||
"XAI_ENROLLER_SOURCE_KIND": "file",
|
||||
"XAI_ENROLLER_SOURCE_FILE": str(tmp_path / "sessions.tsv"),
|
||||
"XAI_ENROLLER_SOURCE_SALT": "test-salt",
|
||||
"XAI_ENROLLER_LEDGER_PATH": str(tmp_path / "ledger.db"),
|
||||
}
|
||||
values.update(overrides)
|
||||
return values
|
||||
|
||||
|
||||
def test_settings_rejects_public_http_cpa_sink(tmp_path):
|
||||
with pytest.raises(ValueError, match="HTTPS"):
|
||||
Settings.from_environ(
|
||||
base_env(
|
||||
tmp_path,
|
||||
XAI_ENROLLER_SINK="cpa",
|
||||
XAI_ENROLLER_CPA_BASE_URL="http://example.test",
|
||||
XAI_ENROLLER_CPA_MANAGEMENT_SECRET="secret",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_settings_defaults_are_bounded_and_redacted(tmp_path):
|
||||
settings = Settings.from_environ(base_env(tmp_path))
|
||||
assert settings.target == 1
|
||||
assert settings.concurrency == 1
|
||||
assert settings.retry_attempts == 0
|
||||
assert settings.executor == "http"
|
||||
assert settings.poll_interval == 5.0
|
||||
redacted = settings.redacted_dict()
|
||||
assert "test-salt" not in repr(redacted)
|
||||
assert "CPA_MANAGEMENT_SECRET" not in repr(redacted)
|
||||
assert "source_salt" not in redacted
|
||||
|
||||
|
||||
def test_settings_accepts_bounded_retry_attempts(tmp_path):
|
||||
settings = Settings.from_environ(
|
||||
base_env(tmp_path, XAI_ENROLLER_RETRY_ATTEMPTS="2")
|
||||
)
|
||||
|
||||
assert settings.retry_attempts == 2
|
||||
|
||||
|
||||
def test_settings_accepts_remote_source_without_a_local_source_file(tmp_path):
|
||||
settings = Settings.from_environ(
|
||||
base_env(
|
||||
tmp_path,
|
||||
XAI_ENROLLER_SOURCE_KIND="remote",
|
||||
XAI_ENROLLER_SOURCE_FILE="",
|
||||
)
|
||||
)
|
||||
|
||||
assert settings.source_kind == "remote"
|
||||
assert settings.source_file is None
|
||||
assert settings.source_db is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("key", "value", "message"),
|
||||
[
|
||||
("XAI_ENROLLER_SOURCE_KIND", "other", "source"),
|
||||
("XAI_ENROLLER_AUTH_EXECUTOR", "other", "executor"),
|
||||
("XAI_ENROLLER_CONCURRENCY", "0", "concurrency"),
|
||||
("XAI_ENROLLER_POLL_SEC", "0", "poll"),
|
||||
("XAI_ENROLLER_RETRY_ATTEMPTS", "-1", "retry"),
|
||||
("XAI_ENROLLER_RETRY_ATTEMPTS", "4", "retry"),
|
||||
("XAI_ENROLLER_TARGET", "0", "target"),
|
||||
("XAI_ENROLLER_TARGET", "101", "target"),
|
||||
],
|
||||
)
|
||||
def test_settings_reject_invalid_bounds(tmp_path, key, value, message):
|
||||
with pytest.raises(ValueError, match=message):
|
||||
Settings.from_environ(base_env(tmp_path, **{key: value}))
|
||||
|
||||
|
||||
def test_file_source_requires_private_regular_file(tmp_path):
|
||||
path = tmp_path / "sessions.tsv"
|
||||
path.write_text("one\tsecret\n", encoding="utf-8")
|
||||
path.chmod(stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP)
|
||||
with pytest.raises(ValueError, match="0600"):
|
||||
list(FileSourceAdapter(path).records())
|
||||
|
||||
|
||||
def test_file_source_suppresses_duplicates_without_leaking_tokens(tmp_path):
|
||||
path = tmp_path / "sessions.tsv"
|
||||
path.write_text("one\tsecret-a\none\tsecret-b\n", encoding="utf-8")
|
||||
path.chmod(0o600)
|
||||
records = list(FileSourceAdapter(path).records())
|
||||
assert len(records) == 1
|
||||
assert records[0].source_id == "one"
|
||||
assert "secret" not in repr(records)
|
||||
|
||||
|
||||
def test_sqlite_source_uses_fixed_read_only_query_and_hmac_ids(tmp_path):
|
||||
path = tmp_path / "accounts.db"
|
||||
connection = sqlite3.connect(path)
|
||||
connection.execute(
|
||||
"CREATE TABLE accounts (token TEXT, status TEXT, deleted_at TEXT, updated_at INTEGER)"
|
||||
)
|
||||
connection.executemany(
|
||||
"INSERT INTO accounts VALUES (?, ?, ?, ?)",
|
||||
[
|
||||
("old", "active", None, 1),
|
||||
("new", "active", None, 2),
|
||||
("deleted", "active", "2026-01-01", 3),
|
||||
("inactive", "inactive", None, 4),
|
||||
],
|
||||
)
|
||||
connection.commit()
|
||||
connection.close()
|
||||
records = list(SQLiteSourceAdapter(path, b"salt").records())
|
||||
assert [record.sso_token for record in records] == ["new", "old"]
|
||||
assert records[0].source_id != records[1].source_id
|
||||
assert "new" not in repr(records)
|
||||
293
tests/test_xai_enroller_coordinator.py
Normal file
293
tests/test_xai_enroller_coordinator.py
Normal file
@@ -0,0 +1,293 @@
|
||||
import asyncio
|
||||
import sqlite3
|
||||
|
||||
import pytest
|
||||
|
||||
from xai_enroller.coordinator import EnrollmentCoordinator
|
||||
from xai_enroller.models import (
|
||||
AuthorizationResult,
|
||||
AuthorizationStatus,
|
||||
DeviceFlow,
|
||||
JobStatus,
|
||||
OAuthCredential,
|
||||
SinkReceipt,
|
||||
SourceRecord,
|
||||
)
|
||||
from xai_enroller.sources import FileSourceAdapter, SQLiteSourceAdapter
|
||||
|
||||
|
||||
class Source:
|
||||
async def records(self):
|
||||
for index in range(3):
|
||||
yield SourceRecord(f"source-{index}", f"token-{index}")
|
||||
|
||||
|
||||
class Protocol:
|
||||
def __init__(self):
|
||||
self.polls = 0
|
||||
|
||||
async def start_device_flow(self):
|
||||
return DeviceFlow("device", "code", "https://accounts.x.ai/device", 60, 0)
|
||||
|
||||
async def poll_token(self, **kwargs):
|
||||
self.polls += 1
|
||||
return OAuthCredential(
|
||||
"access",
|
||||
"refresh",
|
||||
None,
|
||||
"Bearer",
|
||||
3600,
|
||||
"2026-07-11T00:00:00Z",
|
||||
"2026-07-11T00:00:00Z",
|
||||
"subject",
|
||||
"https://auth.x.ai/oauth2/token",
|
||||
)
|
||||
|
||||
|
||||
class Executor:
|
||||
active = 0
|
||||
max_active = 0
|
||||
|
||||
async def confirm(self, source, flow):
|
||||
type(self).active += 1
|
||||
type(self).max_active = max(type(self).max_active, type(self).active)
|
||||
await asyncio.sleep(0)
|
||||
type(self).active -= 1
|
||||
if source.source_id == "source-1":
|
||||
return AuthorizationResult(AuthorizationStatus.NEEDS_BROWSER, "challenge")
|
||||
return AuthorizationResult(AuthorizationStatus.AUTHORIZED, "confirmed")
|
||||
|
||||
|
||||
class Sink:
|
||||
async def store(self, credential):
|
||||
if credential.subject == "subject":
|
||||
return SinkReceipt("receipt")
|
||||
|
||||
|
||||
class FlakyProtocol(Protocol):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.calls = 0
|
||||
|
||||
async def start_device_flow(self):
|
||||
self.calls += 1
|
||||
if self.calls == 1:
|
||||
raise OSError("temporary transport")
|
||||
return await super().start_device_flow()
|
||||
|
||||
|
||||
class InvalidProtocol(Protocol):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.calls = 0
|
||||
|
||||
async def start_device_flow(self):
|
||||
self.calls += 1
|
||||
raise ValueError("invalid device response")
|
||||
|
||||
|
||||
class FailingExecutor(Executor):
|
||||
async def confirm(self, source, flow):
|
||||
raise OSError("confirmation transport")
|
||||
|
||||
|
||||
class SlowDeviceFlowProtocol(Protocol):
|
||||
async def start_device_flow(self):
|
||||
await asyncio.sleep(0.04)
|
||||
return await super().start_device_flow()
|
||||
|
||||
|
||||
class SlowExecutor(Executor):
|
||||
async def confirm(self, source, flow):
|
||||
await asyncio.sleep(0.04)
|
||||
return await super().confirm(source, flow)
|
||||
|
||||
|
||||
def test_coordinator_limits_concurrency_and_maps_terminal_results(tmp_path):
|
||||
protocol = Protocol()
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=Source(),
|
||||
protocol=protocol,
|
||||
executor=Executor(),
|
||||
sink=Sink(),
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"salt",
|
||||
concurrency=2,
|
||||
timeout=5,
|
||||
)
|
||||
results = asyncio.run(coordinator.run(target=3))
|
||||
assert [result.status for result in results] == [
|
||||
JobStatus.IMPORTED,
|
||||
JobStatus.NEEDS_BROWSER,
|
||||
JobStatus.IMPORTED,
|
||||
]
|
||||
assert Executor.max_active == 2
|
||||
assert protocol.polls == 2
|
||||
|
||||
|
||||
def test_coordinator_accepts_explicit_synced_records(tmp_path):
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=Source(),
|
||||
protocol=Protocol(),
|
||||
executor=Executor(),
|
||||
sink=Sink(),
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"salt",
|
||||
concurrency=1,
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
results = asyncio.run(
|
||||
coordinator.run_records([SourceRecord("synced-account", "synced-sso")])
|
||||
)
|
||||
|
||||
assert [result.source_id for result in results] == ["synced-account"]
|
||||
assert results[0].status is JobStatus.IMPORTED
|
||||
|
||||
|
||||
def test_coordinator_emits_device_flow_after_the_flow_is_created(tmp_path):
|
||||
events = []
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=Source(),
|
||||
protocol=Protocol(),
|
||||
executor=Executor(),
|
||||
sink=Sink(),
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"salt",
|
||||
concurrency=1,
|
||||
timeout=5,
|
||||
event_callback=lambda kind, data: events.append((kind, data)),
|
||||
)
|
||||
|
||||
asyncio.run(coordinator.run_records([SourceRecord("synced-account", "synced-sso")]))
|
||||
|
||||
assert events == [
|
||||
(
|
||||
"device_flow",
|
||||
{
|
||||
"source_id": "synced-account",
|
||||
"user_code": "code",
|
||||
"verification_url": "https://accounts.x.ai/device",
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def test_coordinator_without_sink_never_reports_imported(tmp_path):
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=Source(),
|
||||
protocol=Protocol(),
|
||||
executor=Executor(),
|
||||
sink=None,
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"salt",
|
||||
concurrency=1,
|
||||
timeout=5,
|
||||
)
|
||||
results = asyncio.run(coordinator.run(target=1))
|
||||
assert results[0].status is JobStatus.SINK_FAILED
|
||||
|
||||
|
||||
def test_coordinator_retries_only_before_device_flow_creation(tmp_path):
|
||||
protocol = FlakyProtocol()
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=Source(),
|
||||
protocol=protocol,
|
||||
executor=Executor(),
|
||||
sink=Sink(),
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"salt",
|
||||
concurrency=1,
|
||||
timeout=5,
|
||||
retry_attempts=1,
|
||||
)
|
||||
results = asyncio.run(coordinator.run(target=1))
|
||||
assert results[0].status is JobStatus.IMPORTED
|
||||
assert protocol.calls == 2
|
||||
|
||||
|
||||
def test_coordinator_does_not_retry_after_device_flow_creation(tmp_path):
|
||||
protocol = Protocol()
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=Source(),
|
||||
protocol=protocol,
|
||||
executor=FailingExecutor(),
|
||||
sink=Sink(),
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"salt",
|
||||
concurrency=1,
|
||||
timeout=5,
|
||||
retry_attempts=3,
|
||||
)
|
||||
results = asyncio.run(coordinator.run(target=1))
|
||||
assert results[0].status is JobStatus.TRANSPORT_FAILED
|
||||
assert protocol.polls == 0
|
||||
|
||||
|
||||
def test_coordinator_does_not_retry_non_transport_device_flow_error(tmp_path):
|
||||
protocol = InvalidProtocol()
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=Source(),
|
||||
protocol=protocol,
|
||||
executor=Executor(),
|
||||
sink=Sink(),
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"salt",
|
||||
concurrency=1,
|
||||
timeout=5,
|
||||
retry_attempts=3,
|
||||
)
|
||||
results = asyncio.run(coordinator.run(target=1))
|
||||
assert results[0].status is JobStatus.TRANSPORT_FAILED
|
||||
assert protocol.calls == 1
|
||||
|
||||
|
||||
def test_coordinator_applies_timeout_to_the_entire_attempt(tmp_path):
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=Source(),
|
||||
protocol=SlowDeviceFlowProtocol(),
|
||||
executor=SlowExecutor(),
|
||||
sink=Sink(),
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"salt",
|
||||
concurrency=1,
|
||||
timeout=0.06,
|
||||
retry_attempts=3,
|
||||
)
|
||||
|
||||
results = asyncio.run(coordinator.run(target=1))
|
||||
|
||||
assert results[0].status is JobStatus.TIMEOUT
|
||||
|
||||
|
||||
@pytest.mark.parametrize("source_kind", ["file", "sqlite"])
|
||||
def test_coordinator_consumes_real_source_adapters(tmp_path, source_kind):
|
||||
if source_kind == "file":
|
||||
source_path = tmp_path / "sessions.tsv"
|
||||
source_path.write_text("file-source\tfile-token\n", encoding="utf-8")
|
||||
source_path.chmod(0o600)
|
||||
source = FileSourceAdapter(source_path)
|
||||
else:
|
||||
source_path = tmp_path / "accounts.db"
|
||||
connection = sqlite3.connect(source_path)
|
||||
connection.execute(
|
||||
"CREATE TABLE accounts (token TEXT, status TEXT, deleted_at TEXT, updated_at INTEGER)"
|
||||
)
|
||||
connection.execute("INSERT INTO accounts VALUES (?, ?, ?, ?)", ("db-token", "active", None, 1))
|
||||
connection.commit()
|
||||
connection.close()
|
||||
source = SQLiteSourceAdapter(source_path, b"source-salt")
|
||||
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=source,
|
||||
protocol=Protocol(),
|
||||
executor=Executor(),
|
||||
sink=Sink(),
|
||||
ledger_path=tmp_path / f"{source_kind}-ledger.db",
|
||||
ledger_salt=b"ledger-salt",
|
||||
concurrency=1,
|
||||
timeout=5,
|
||||
)
|
||||
results = asyncio.run(coordinator.run(target=1))
|
||||
assert len(results) == 1
|
||||
assert results[0].status is JobStatus.IMPORTED
|
||||
411
tests/test_xai_enroller_executors.py
Normal file
411
tests/test_xai_enroller_executors.py
Normal file
@@ -0,0 +1,411 @@
|
||||
import asyncio
|
||||
from contextlib import suppress
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from xai_enroller.executors import HTTPProbeExecutor, PlaywrightExecutor
|
||||
from xai_enroller.models import AuthorizationStatus, DeviceFlow, SourceRecord
|
||||
|
||||
|
||||
def test_http_probe_classifies_403_challenge_without_submission():
|
||||
calls = []
|
||||
|
||||
def handler(request):
|
||||
calls.append(request)
|
||||
return httpx.Response(403, text="<html>challenge-platform</html>")
|
||||
|
||||
executor = HTTPProbeExecutor(
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler), follow_redirects=False)
|
||||
)
|
||||
source = SourceRecord("source", "sso-token")
|
||||
flow = DeviceFlow("device", "ABCD", "https://accounts.x.ai/oauth2/device?user_code=ABCD", 60, 1)
|
||||
result = asyncio.run(executor.confirm(source, flow))
|
||||
assert result.status is AuthorizationStatus.NEEDS_BROWSER
|
||||
assert len(calls) == 1
|
||||
assert calls[0].method == "GET"
|
||||
assert "sso-token" not in repr(result)
|
||||
|
||||
|
||||
def test_http_probe_does_not_follow_redirects_or_submit_forms():
|
||||
def handler(request):
|
||||
return httpx.Response(302, headers={"location": "https://evil.example/"})
|
||||
|
||||
executor = HTTPProbeExecutor(
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler), follow_redirects=False)
|
||||
)
|
||||
result = asyncio.run(
|
||||
executor.confirm(
|
||||
SourceRecord("source", "token"),
|
||||
DeviceFlow("device", "ABCD", "https://accounts.x.ai/oauth2/device", 60, 1),
|
||||
)
|
||||
)
|
||||
assert result.status is AuthorizationStatus.NEEDS_INTERACTION
|
||||
|
||||
|
||||
class FakeButton:
|
||||
def __init__(self, count):
|
||||
self._count = count
|
||||
|
||||
async def count(self):
|
||||
return self._count
|
||||
|
||||
@property
|
||||
def first(self):
|
||||
return self
|
||||
|
||||
async def click(self):
|
||||
return None
|
||||
|
||||
|
||||
class FakePage:
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
|
||||
async def goto(self, url, wait_until):
|
||||
self.url = url
|
||||
|
||||
def locator(self, selector):
|
||||
return self
|
||||
|
||||
async def inner_text(self):
|
||||
return "Authorize access"
|
||||
|
||||
def get_by_role(self, role, name):
|
||||
return FakeButton(1 if name == "Authorize" else 0)
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class FakeContext:
|
||||
def __init__(self):
|
||||
self.cookies = []
|
||||
self.page = FakePage()
|
||||
self.closed = False
|
||||
|
||||
async def add_cookies(self, cookies):
|
||||
self.cookies.extend(cookies)
|
||||
|
||||
async def new_page(self):
|
||||
return self.page
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class FakeBrowser:
|
||||
def __init__(self):
|
||||
self.launches = 0
|
||||
self.contexts = []
|
||||
self.closed = False
|
||||
|
||||
async def new_context(self):
|
||||
context = FakeContext()
|
||||
self.contexts.append(context)
|
||||
return context
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class FakeChromium:
|
||||
def __init__(self, browser):
|
||||
self.browser = browser
|
||||
self.launch_options = None
|
||||
|
||||
async def launch(self, **options):
|
||||
self.launch_options = options
|
||||
self.browser.launches += 1
|
||||
return self.browser
|
||||
|
||||
|
||||
class FakePlaywright:
|
||||
def __init__(self, browser):
|
||||
self.chromium = FakeChromium(browser)
|
||||
self.stopped = False
|
||||
|
||||
async def stop(self):
|
||||
self.stopped = True
|
||||
|
||||
|
||||
class FakePlaywrightFactory:
|
||||
def __init__(self):
|
||||
self.browser = FakeBrowser()
|
||||
self.playwright = FakePlaywright(self.browser)
|
||||
|
||||
def __call__(self):
|
||||
factory = self
|
||||
|
||||
class Runner:
|
||||
async def start(self):
|
||||
return factory.playwright
|
||||
|
||||
return Runner()
|
||||
|
||||
|
||||
class FailingPageContext(FakeContext):
|
||||
async def new_page(self):
|
||||
raise RuntimeError("page creation failed")
|
||||
|
||||
|
||||
class FailingPageBrowser(FakeBrowser):
|
||||
async def new_context(self):
|
||||
context = FailingPageContext()
|
||||
self.contexts.append(context)
|
||||
return context
|
||||
|
||||
|
||||
def test_playwright_executor_isolates_context_and_scopes_exact_cookie():
|
||||
factory = FakePlaywrightFactory()
|
||||
executor = PlaywrightExecutor(playwright_factory=factory)
|
||||
source = SourceRecord("source", "secret-token")
|
||||
flow = DeviceFlow("device", "ABCD", "https://accounts.x.ai/oauth2/device", 60, 1)
|
||||
result = asyncio.run(executor.confirm(source, flow))
|
||||
asyncio.run(executor.close())
|
||||
assert result.status is AuthorizationStatus.AUTHORIZED
|
||||
assert factory.browser.launches == 1
|
||||
context = factory.browser.contexts[0]
|
||||
assert context.cookies == [{
|
||||
"name": "sso",
|
||||
"value": "secret-token",
|
||||
"domain": "accounts.x.ai",
|
||||
"path": "/",
|
||||
"secure": True,
|
||||
"httpOnly": True,
|
||||
}]
|
||||
assert context.page.closed
|
||||
assert context.closed
|
||||
assert factory.browser.closed
|
||||
|
||||
|
||||
def test_playwright_executor_closes_context_when_page_creation_fails():
|
||||
factory = FakePlaywrightFactory()
|
||||
factory.browser = FailingPageBrowser()
|
||||
factory.playwright = FakePlaywright(factory.browser)
|
||||
executor = PlaywrightExecutor(playwright_factory=factory)
|
||||
|
||||
result = asyncio.run(
|
||||
executor.confirm(
|
||||
SourceRecord("source", "secret-token"),
|
||||
DeviceFlow("device", "ABCD", "https://accounts.x.ai/oauth2/device", 60, 1),
|
||||
)
|
||||
)
|
||||
asyncio.run(executor.close())
|
||||
|
||||
assert result.status is AuthorizationStatus.NEEDS_INTERACTION
|
||||
assert factory.browser.contexts[0].closed
|
||||
|
||||
|
||||
def test_playwright_executor_launches_configured_fingerprint_browser():
|
||||
factory = FakePlaywrightFactory()
|
||||
executor = PlaywrightExecutor(
|
||||
playwright_factory=factory,
|
||||
executable_path="/opt/cloakbrowser/chrome",
|
||||
)
|
||||
|
||||
asyncio.run(executor.start())
|
||||
asyncio.run(executor.close())
|
||||
|
||||
assert factory.playwright.chromium.launch_options == {
|
||||
"headless": True,
|
||||
"executable_path": "/opt/cloakbrowser/chrome",
|
||||
}
|
||||
|
||||
|
||||
def test_playwright_executor_uses_fingerprint_browser_path_from_environment(monkeypatch):
|
||||
monkeypatch.setenv("XAI_ENROLLER_BROWSER_EXECUTABLE", "/opt/cloakbrowser/chrome")
|
||||
factory = FakePlaywrightFactory()
|
||||
executor = PlaywrightExecutor(playwright_factory=factory)
|
||||
|
||||
asyncio.run(executor.start())
|
||||
asyncio.run(executor.close())
|
||||
|
||||
assert factory.playwright.chromium.launch_options == {
|
||||
"headless": True,
|
||||
"executable_path": "/opt/cloakbrowser/chrome",
|
||||
}
|
||||
|
||||
|
||||
def test_playwright_executor_finds_macos_cloakbrowser_binary(monkeypatch):
|
||||
macos_binary = (
|
||||
"/Users/test/.cloakbrowser/chromium-145/Chromium.app/Contents/MacOS/Chromium"
|
||||
)
|
||||
monkeypatch.delenv("XAI_ENROLLER_BROWSER_EXECUTABLE", raising=False)
|
||||
monkeypatch.setattr(
|
||||
"xai_enroller.executors.glob.glob",
|
||||
lambda pattern: [macos_binary] if "Chromium.app" in pattern else [],
|
||||
)
|
||||
|
||||
assert PlaywrightExecutor._find_executable_path() == macos_binary
|
||||
|
||||
|
||||
def test_playwright_attempt_cleanup_defers_repeated_cancellation():
|
||||
class Page:
|
||||
def __init__(self):
|
||||
self.entered = asyncio.Event()
|
||||
self.release = asyncio.Event()
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
self.entered.set()
|
||||
await self.release.wait()
|
||||
self.closed = True
|
||||
|
||||
class Context:
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
async def scenario():
|
||||
page = Page()
|
||||
context = Context()
|
||||
task = asyncio.create_task(
|
||||
PlaywrightExecutor._close_attempt_resources(page, context)
|
||||
)
|
||||
await page.entered.wait()
|
||||
task.cancel()
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
page.release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert page.closed
|
||||
assert context.closed
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_playwright_attempt_cleanup_has_a_true_wall_clock_deadline(monkeypatch):
|
||||
class Page:
|
||||
def __init__(self):
|
||||
self.entered = asyncio.Event()
|
||||
self.release = asyncio.Event()
|
||||
|
||||
async def close(self):
|
||||
self.entered.set()
|
||||
while not self.release.is_set():
|
||||
try:
|
||||
await self.release.wait()
|
||||
except asyncio.CancelledError:
|
||||
# Playwright transports can remain stuck while cancellation
|
||||
# propagates. A hard deadline must not await that forever.
|
||||
continue
|
||||
|
||||
async def scenario():
|
||||
monkeypatch.setattr(PlaywrightExecutor, "CLOSE_TIMEOUT_SECONDS", 0.02)
|
||||
page = Page()
|
||||
task = asyncio.create_task(
|
||||
PlaywrightExecutor._close_attempt_resources(page, None)
|
||||
)
|
||||
await page.entered.wait()
|
||||
done, _pending = await asyncio.wait({task}, timeout=0.10)
|
||||
completed_within_deadline = task in done
|
||||
|
||||
page.release.set()
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
assert completed_within_deadline
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_playwright_confirmation_has_hard_deadline_and_restarts_browser(monkeypatch):
|
||||
class StubbornGotoPage(FakePage):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.entered = asyncio.Event()
|
||||
self.release = asyncio.Event()
|
||||
|
||||
async def goto(self, url, wait_until):
|
||||
self.url = url
|
||||
self.entered.set()
|
||||
while not self.release.is_set():
|
||||
try:
|
||||
await self.release.wait()
|
||||
except asyncio.CancelledError:
|
||||
# Model a wedged Playwright transport that does not finish
|
||||
# propagating task cancellation.
|
||||
continue
|
||||
|
||||
class StubbornGotoContext(FakeContext):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.page = StubbornGotoPage()
|
||||
|
||||
class StubbornGotoBrowser(FakeBrowser):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.close_entered = asyncio.Event()
|
||||
self.close_release = asyncio.Event()
|
||||
|
||||
async def new_context(self):
|
||||
context = StubbornGotoContext()
|
||||
self.contexts.append(context)
|
||||
return context
|
||||
|
||||
async def close(self):
|
||||
self.close_entered.set()
|
||||
while not self.close_release.is_set():
|
||||
try:
|
||||
await self.close_release.wait()
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
self.closed = True
|
||||
|
||||
class RotatingPlaywrightFactory:
|
||||
def __init__(self):
|
||||
self.browsers = [StubbornGotoBrowser(), FakeBrowser()]
|
||||
self.playwrights = [FakePlaywright(browser) for browser in self.browsers]
|
||||
self.starts = 0
|
||||
|
||||
def __call__(self):
|
||||
factory = self
|
||||
|
||||
class Runner:
|
||||
async def start(self):
|
||||
playwright = factory.playwrights[factory.starts]
|
||||
factory.starts += 1
|
||||
return playwright
|
||||
|
||||
return Runner()
|
||||
|
||||
async def scenario():
|
||||
monkeypatch.setattr(PlaywrightExecutor, "ATTEMPT_TIMEOUT_SECONDS", 0.02)
|
||||
monkeypatch.setattr(PlaywrightExecutor, "CLOSE_TIMEOUT_SECONDS", 0.02)
|
||||
factory = RotatingPlaywrightFactory()
|
||||
executor = PlaywrightExecutor(playwright_factory=factory)
|
||||
source = SourceRecord("source", "secret-token")
|
||||
flow = DeviceFlow(
|
||||
"device", "ABCD", "https://accounts.x.ai/oauth2/device", 60, 1
|
||||
)
|
||||
|
||||
first = asyncio.create_task(executor.confirm(source, flow))
|
||||
while not factory.browsers[0].contexts:
|
||||
await asyncio.sleep(0)
|
||||
stuck_page = factory.browsers[0].contexts[0].page
|
||||
await stuck_page.entered.wait()
|
||||
done, _pending = await asyncio.wait({first}, timeout=0.10)
|
||||
completed_within_deadline = first in done
|
||||
|
||||
stuck_page.release.set()
|
||||
factory.browsers[0].close_release.set()
|
||||
first_result = await first
|
||||
await asyncio.sleep(0)
|
||||
second_result = await executor.confirm(source, flow)
|
||||
await executor.close()
|
||||
|
||||
assert completed_within_deadline
|
||||
assert first_result.status is AuthorizationStatus.NEEDS_INTERACTION
|
||||
assert first_result.reason_code == "confirmation_timeout"
|
||||
assert second_result.status is AuthorizationStatus.AUTHORIZED
|
||||
assert factory.starts == 2
|
||||
assert factory.browsers[0].close_entered.is_set()
|
||||
assert factory.browsers[0].closed
|
||||
assert factory.playwrights[0].stopped
|
||||
|
||||
asyncio.run(scenario())
|
||||
188
tests/test_xai_enroller_integration.py
Normal file
188
tests/test_xai_enroller_integration.py
Normal file
@@ -0,0 +1,188 @@
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import httpx
|
||||
|
||||
from xai_enroller.coordinator import EnrollmentCoordinator
|
||||
from xai_enroller.executors import PlaywrightExecutor
|
||||
from xai_enroller.models import SourceRecord
|
||||
from xai_enroller.protocol import XAIProfile, XAIProtocol
|
||||
from xai_enroller.sinks import CPAAuthFileSink
|
||||
|
||||
|
||||
class SingleSource:
|
||||
async def records(self):
|
||||
yield SourceRecord("source-1", "sso-secret")
|
||||
|
||||
|
||||
class FakeButton:
|
||||
async def count(self):
|
||||
return 1
|
||||
|
||||
@property
|
||||
def first(self):
|
||||
return self
|
||||
|
||||
async def click(self):
|
||||
return None
|
||||
|
||||
|
||||
class FakePage:
|
||||
async def goto(self, url, wait_until):
|
||||
self.url = url
|
||||
|
||||
def locator(self, selector):
|
||||
return self
|
||||
|
||||
async def inner_text(self):
|
||||
return "Authorize access"
|
||||
|
||||
def get_by_role(self, role, name):
|
||||
return FakeButton() if name == "Authorize" else type(
|
||||
"NoButton",
|
||||
(),
|
||||
{"count": staticmethod(lambda: asyncio.sleep(0, result=0))},
|
||||
)()
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
|
||||
class FakeContext:
|
||||
def __init__(self, browser):
|
||||
self.browser = browser
|
||||
self.cookies = []
|
||||
|
||||
async def add_cookies(self, cookies):
|
||||
self.cookies.extend(cookies)
|
||||
|
||||
async def new_page(self):
|
||||
return FakePage()
|
||||
|
||||
async def close(self):
|
||||
self.browser.live_contexts -= 1
|
||||
|
||||
|
||||
class FakeBrowser:
|
||||
def __init__(self):
|
||||
self.launches = 0
|
||||
self.live_contexts = 0
|
||||
self.max_live_contexts = 0
|
||||
|
||||
async def new_context(self):
|
||||
self.live_contexts += 1
|
||||
self.max_live_contexts = max(self.max_live_contexts, self.live_contexts)
|
||||
return FakeContext(self)
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
|
||||
class FakeChromium:
|
||||
def __init__(self, browser):
|
||||
self.browser = browser
|
||||
|
||||
async def launch(self, **options):
|
||||
self.browser.launches += 1
|
||||
return self.browser
|
||||
|
||||
|
||||
class FakePlaywright:
|
||||
def __init__(self, browser):
|
||||
self.chromium = FakeChromium(browser)
|
||||
|
||||
async def stop(self):
|
||||
return None
|
||||
|
||||
|
||||
class FakePlaywrightFactory:
|
||||
def __init__(self):
|
||||
self.browser = FakeBrowser()
|
||||
self.playwright = FakePlaywright(self.browser)
|
||||
|
||||
def __call__(self):
|
||||
factory = self
|
||||
|
||||
class Runner:
|
||||
async def start(self):
|
||||
return factory.playwright
|
||||
|
||||
return Runner()
|
||||
|
||||
|
||||
def test_adapter_boundary_hands_xai_credential_to_cpa_without_ledger_secrets(tmp_path):
|
||||
def xai_handler(request):
|
||||
if request.url.path == "/.well-known/openid-configuration":
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_authorization_endpoint": "https://auth.x.ai/oauth2/device/code",
|
||||
"token_endpoint": "https://auth.x.ai/oauth2/token",
|
||||
},
|
||||
)
|
||||
if request.url.path == "/oauth2/device/code":
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_code": "device-secret",
|
||||
"user_code": "ABCD",
|
||||
"verification_uri": "https://accounts.x.ai/oauth2/device",
|
||||
"verification_uri_complete": (
|
||||
"https://accounts.x.ai/oauth2/device?user_code=ABCD"
|
||||
),
|
||||
"expires_in": 60,
|
||||
"interval": 0,
|
||||
},
|
||||
)
|
||||
if request.url.path == "/oauth2/token":
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"access_token": "access-secret",
|
||||
"refresh_token": "refresh-secret",
|
||||
"id_token": "id-secret",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600,
|
||||
"sub": "subject-1",
|
||||
},
|
||||
)
|
||||
raise AssertionError(f"unexpected xAI request: {request.url}")
|
||||
|
||||
cpa_requests = []
|
||||
|
||||
def cpa_handler(request):
|
||||
cpa_requests.append(request)
|
||||
return httpx.Response(201)
|
||||
|
||||
xai_client = httpx.AsyncClient(transport=httpx.MockTransport(xai_handler))
|
||||
cpa_client = httpx.AsyncClient(transport=httpx.MockTransport(cpa_handler))
|
||||
factory = FakePlaywrightFactory()
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=SingleSource(),
|
||||
protocol=XAIProtocol(xai_client, XAIProfile.default()),
|
||||
executor=PlaywrightExecutor(playwright_factory=factory),
|
||||
sink=CPAAuthFileSink(
|
||||
"https://cpa.example",
|
||||
"management-secret",
|
||||
cpa_client,
|
||||
name_secret=b"name-secret",
|
||||
),
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"ledger-salt",
|
||||
)
|
||||
|
||||
try:
|
||||
results = asyncio.run(coordinator.run(target=1))
|
||||
finally:
|
||||
asyncio.run(xai_client.aclose())
|
||||
asyncio.run(cpa_client.aclose())
|
||||
|
||||
assert results[0].status.value == "imported"
|
||||
assert factory.browser.launches == 1
|
||||
assert factory.browser.max_live_contexts == 1
|
||||
assert cpa_requests[0].url.path == "/v0/management/auth-files"
|
||||
assert cpa_requests[0].url.params["name"].startswith("xai-")
|
||||
assert json.loads(cpa_requests[0].content)["refresh_token"] == "refresh-secret"
|
||||
ledger_bytes = (tmp_path / "ledger.db").read_bytes()
|
||||
for secret in ("sso-secret", "device-secret", "access-secret", "refresh-secret", "id-secret"):
|
||||
assert secret.encode() not in ledger_bytes
|
||||
163
tests/test_xai_enroller_ledger.py
Normal file
163
tests/test_xai_enroller_ledger.py
Normal file
@@ -0,0 +1,163 @@
|
||||
import sqlite3
|
||||
|
||||
import pytest
|
||||
|
||||
from xai_enroller.ledger import Ledger
|
||||
from xai_enroller.models import JobStatus
|
||||
|
||||
|
||||
def test_ledger_persists_only_redacted_terminal_fields(tmp_path):
|
||||
path = tmp_path / "ledger.db"
|
||||
ledger = Ledger(path, b"salt")
|
||||
job_id = ledger.start("source", attempt=1)
|
||||
ledger.finish(job_id, JobStatus.SINK_FAILED, "sink_failed", "receipt")
|
||||
raw = path.read_bytes()
|
||||
for secret in [
|
||||
"sso-token",
|
||||
"device-code",
|
||||
"https://accounts.x.ai",
|
||||
"access-token",
|
||||
"refresh-token",
|
||||
"id-token",
|
||||
"person@example.com",
|
||||
]:
|
||||
assert secret.encode() not in raw
|
||||
row = ledger.get(job_id)
|
||||
assert row["status"] == "sink_failed"
|
||||
assert row["reason_code"] == "sink_failed"
|
||||
assert row["sink_receipt_fingerprint"] == "receipt"
|
||||
assert "source" not in repr(row["source_fingerprint"])
|
||||
|
||||
|
||||
def test_ledger_recovers_pending_jobs(tmp_path):
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
job_id = ledger.start("source", attempt=1)
|
||||
ledger.recover_pending()
|
||||
assert ledger.get(job_id)["status"] == JobStatus.CANCELLED.value
|
||||
|
||||
|
||||
def test_ledger_backfills_imported_receipts_into_available_inventory(tmp_path):
|
||||
path = tmp_path / "ledger.db"
|
||||
with sqlite3.connect(path) as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
CREATE TABLE jobs (
|
||||
job_id INTEGER PRIMARY KEY,
|
||||
source_fingerprint TEXT NOT NULL,
|
||||
attempt_number INTEGER NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
started_at TEXT NOT NULL,
|
||||
finished_at TEXT,
|
||||
reason_code TEXT,
|
||||
sink_receipt_fingerprint TEXT
|
||||
)
|
||||
"""
|
||||
)
|
||||
connection.executemany(
|
||||
"INSERT INTO jobs(source_fingerprint, attempt_number, status, started_at, "
|
||||
"finished_at, reason_code, sink_receipt_fingerprint) VALUES (?, 1, ?, ?, ?, ?, ?)",
|
||||
[
|
||||
("source-old", "imported", "2026-01-01", "2026-01-02", "imported", "receipt-old"),
|
||||
("source-new", "imported", "2026-01-03", "2026-01-04", "imported", "receipt-new"),
|
||||
("source-failed", "sink_failed", "2026-01-05", "2026-01-06", "sink_failed", "receipt-failed"),
|
||||
],
|
||||
)
|
||||
|
||||
ledger = Ledger(path, b"salt")
|
||||
|
||||
assert ledger.inventory_counts() == {
|
||||
"available": 2,
|
||||
"claiming": 0,
|
||||
"claimed": 0,
|
||||
}
|
||||
assert ledger.claim_available(2, "batch-backfill") == [
|
||||
"receipt-new",
|
||||
"receipt-old",
|
||||
]
|
||||
|
||||
|
||||
def test_imported_finish_and_inventory_insert_are_atomic(tmp_path):
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
job_id = ledger.start("source", attempt=1)
|
||||
with sqlite3.connect(ledger.path) as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
CREATE TRIGGER reject_inventory_insert
|
||||
BEFORE INSERT ON credential_inventory
|
||||
BEGIN
|
||||
SELECT RAISE(ABORT, 'inventory insert rejected');
|
||||
END
|
||||
"""
|
||||
)
|
||||
|
||||
with pytest.raises(sqlite3.IntegrityError, match="inventory insert rejected"):
|
||||
ledger.finish(job_id, JobStatus.IMPORTED, "imported", "receipt")
|
||||
|
||||
assert ledger.get(job_id)["status"] == "pending"
|
||||
assert ledger.inventory_counts()["available"] == 0
|
||||
|
||||
|
||||
def test_claim_available_uses_latest_finish_and_tracks_batch(tmp_path):
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
jobs = []
|
||||
for source, receipt in [
|
||||
("source-a", "receipt-a"),
|
||||
("source-b", "receipt-b"),
|
||||
("source-c", "receipt-c"),
|
||||
]:
|
||||
job_id = ledger.start(source)
|
||||
ledger.finish(job_id, JobStatus.IMPORTED, "imported", receipt)
|
||||
jobs.append(job_id)
|
||||
with sqlite3.connect(ledger.path) as connection:
|
||||
connection.executemany(
|
||||
"UPDATE jobs SET finished_at=? WHERE job_id=?",
|
||||
[
|
||||
("2026-01-01T00:00:00+00:00", jobs[0]),
|
||||
("2026-01-03T00:00:00+00:00", jobs[1]),
|
||||
("2026-01-02T00:00:00+00:00", jobs[2]),
|
||||
],
|
||||
)
|
||||
|
||||
assert ledger.claim_available(2, "batch-1") == ["receipt-b", "receipt-c"]
|
||||
assert ledger.claim_available(2, "batch-2") == ["receipt-a"]
|
||||
assert ledger.pending_claims("batch-1") == [
|
||||
{
|
||||
"sink_receipt_fingerprint": "receipt-b",
|
||||
"batch_id": "batch-1",
|
||||
"note": "",
|
||||
},
|
||||
{
|
||||
"sink_receipt_fingerprint": "receipt-c",
|
||||
"batch_id": "batch-1",
|
||||
"note": "",
|
||||
},
|
||||
]
|
||||
assert ledger.inventory_counts() == {
|
||||
"available": 0,
|
||||
"claiming": 3,
|
||||
"claimed": 0,
|
||||
}
|
||||
|
||||
|
||||
def test_claim_completion_and_recovery_are_batch_scoped(tmp_path):
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
for source, receipt in [
|
||||
("source-a", "receipt-a"),
|
||||
("source-b", "receipt-b"),
|
||||
("source-c", "receipt-c"),
|
||||
]:
|
||||
job_id = ledger.start(source)
|
||||
ledger.finish(job_id, JobStatus.IMPORTED, "imported", receipt)
|
||||
|
||||
assert ledger.claim_available(2, "batch-complete")
|
||||
assert ledger.claim_available(1, "batch-recover")
|
||||
assert ledger.mark_claimed("batch-complete", "delivered") == 2
|
||||
assert ledger.recover_claiming("batch-recover", note="worker_restart") == 1
|
||||
|
||||
assert ledger.pending_claims() == []
|
||||
assert ledger.inventory_counts() == {
|
||||
"available": 1,
|
||||
"claiming": 0,
|
||||
"claimed": 2,
|
||||
}
|
||||
assert ledger.claim_available(1, "batch-retry") == ["receipt-a"]
|
||||
245
tests/test_xai_enroller_protocol.py
Normal file
245
tests/test_xai_enroller_protocol.py
Normal file
@@ -0,0 +1,245 @@
|
||||
import base64
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from xai_enroller.models import AuthorizationStatus
|
||||
from xai_enroller.protocol import XAIProtocol, XAIProfile
|
||||
|
||||
|
||||
def _unsigned_jwt(payload):
|
||||
encoded = base64.urlsafe_b64encode(
|
||||
json.dumps(payload, separators=(",", ":")).encode()
|
||||
).decode().rstrip("=")
|
||||
return f"header.{encoded}.signature"
|
||||
|
||||
|
||||
def test_default_profile_matches_observed_grok_cli_device_contract():
|
||||
profile = XAIProfile.default()
|
||||
|
||||
assert profile.client_id == "b1a00492-073a-47ea-816f-4c329264a828"
|
||||
assert profile.scope == "openid profile email offline_access grok-cli:access api:access"
|
||||
assert profile.version == "grok-cli-device-v1"
|
||||
|
||||
|
||||
def test_device_flow_uses_discovered_endpoints_and_profile_scope():
|
||||
requests = []
|
||||
|
||||
def handler(request):
|
||||
requests.append(request)
|
||||
if request.url.path == "/.well-known/openid-configuration":
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_authorization_endpoint": "https://auth.x.ai/oauth2/device/code",
|
||||
"token_endpoint": "https://auth.x.ai/oauth2/token",
|
||||
},
|
||||
)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_code": "device",
|
||||
"user_code": "ABCD",
|
||||
"verification_uri": "https://accounts.x.ai/oauth2/device",
|
||||
"expires_in": 60,
|
||||
"interval": 1,
|
||||
},
|
||||
)
|
||||
|
||||
profile = XAIProfile(client_id="client", scope="openid offline_access")
|
||||
protocol = XAIProtocol(httpx.AsyncClient(transport=httpx.MockTransport(handler)), profile)
|
||||
|
||||
async def run():
|
||||
flow = await protocol.start_device_flow()
|
||||
return flow
|
||||
|
||||
flow = __import__("asyncio").run(run())
|
||||
assert flow.verification_url == "https://accounts.x.ai/oauth2/device?user_code=ABCD"
|
||||
assert dict(__import__("urllib.parse", fromlist=["parse_qsl"]).parse_qsl(requests[1].content.decode())) == {
|
||||
"client_id": "client",
|
||||
"scope": "openid offline_access",
|
||||
}
|
||||
|
||||
|
||||
def test_device_flow_uses_configured_poll_interval_when_issuer_omits_it():
|
||||
def handler(request):
|
||||
if request.url.path == "/.well-known/openid-configuration":
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_authorization_endpoint": "https://auth.x.ai/oauth2/device/code",
|
||||
"token_endpoint": "https://auth.x.ai/oauth2/token",
|
||||
},
|
||||
)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_code": "device",
|
||||
"user_code": "ABCD",
|
||||
"verification_uri": "https://accounts.x.ai/oauth2/device",
|
||||
"expires_in": 60,
|
||||
},
|
||||
)
|
||||
|
||||
protocol = XAIProtocol(
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler)),
|
||||
XAIProfile(client_id="client", scope="openid"),
|
||||
default_poll_interval=7,
|
||||
)
|
||||
|
||||
flow = __import__("asyncio").run(protocol.start_device_flow())
|
||||
|
||||
assert flow.interval == 7
|
||||
|
||||
|
||||
def test_device_flow_preserves_complete_verification_url():
|
||||
def handler(request):
|
||||
if request.url.path == "/.well-known/openid-configuration":
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_authorization_endpoint": "https://auth.x.ai/oauth2/device/code",
|
||||
"token_endpoint": "https://auth.x.ai/oauth2/token",
|
||||
},
|
||||
)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_code": "device",
|
||||
"user_code": "ABCD",
|
||||
"verification_uri": "https://accounts.x.ai/oauth2/device",
|
||||
"verification_uri_complete": "https://accounts.x.ai/oauth2/device?user_code=ABCD&state=issuer",
|
||||
"expires_in": 60,
|
||||
},
|
||||
)
|
||||
|
||||
protocol = XAIProtocol(
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler)),
|
||||
XAIProfile(client_id="client", scope="openid"),
|
||||
)
|
||||
|
||||
flow = __import__("asyncio").run(protocol.start_device_flow())
|
||||
|
||||
assert flow.verification_url == (
|
||||
"https://accounts.x.ai/oauth2/device?user_code=ABCD&state=issuer"
|
||||
)
|
||||
|
||||
|
||||
def test_xai_url_allowlist_rejects_missing_hostname():
|
||||
assert XAIProtocol._allowed_url("https:///oauth2/device") is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"discovery",
|
||||
[
|
||||
{
|
||||
"device_authorization_endpoint": "http://auth.x.ai/device",
|
||||
"token_endpoint": "https://auth.x.ai/token",
|
||||
},
|
||||
{
|
||||
"device_authorization_endpoint": "https://evil.example/device",
|
||||
"token_endpoint": "https://auth.x.ai/token",
|
||||
},
|
||||
],
|
||||
)
|
||||
def test_discovery_rejects_non_xai_or_non_https(discovery):
|
||||
def handler(request):
|
||||
return httpx.Response(200, json=discovery)
|
||||
|
||||
protocol = XAIProtocol(
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler)),
|
||||
XAIProfile(client_id="client", scope="openid"),
|
||||
)
|
||||
with pytest.raises(ValueError, match="allowlist"):
|
||||
__import__("asyncio").run(protocol.discover())
|
||||
|
||||
|
||||
def test_token_polling_maps_pending_slow_down_and_denial():
|
||||
responses = iter(
|
||||
[
|
||||
httpx.Response(400, json={"error": "authorization_pending"}),
|
||||
httpx.Response(400, json={"error": "slow_down"}),
|
||||
httpx.Response(400, json={"error": "access_denied"}),
|
||||
]
|
||||
)
|
||||
|
||||
def handler(request):
|
||||
return next(responses)
|
||||
|
||||
protocol = XAIProtocol(
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler)),
|
||||
XAIProfile(client_id="client", scope="openid"),
|
||||
sleep=lambda _: __import__("asyncio").sleep(0),
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="oauth_denied"):
|
||||
__import__("asyncio").run(
|
||||
protocol.poll_token(
|
||||
endpoint="https://auth.x.ai/oauth2/token",
|
||||
flow=type("Flow", (), {"device_code": "device", "interval": 0})(),
|
||||
timeout=1,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_token_response_requires_refresh_token_and_never_exposes_body():
|
||||
def handler(request):
|
||||
return httpx.Response(200, json={"access_token": "access", "token_type": "Bearer"})
|
||||
|
||||
protocol = XAIProtocol(
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler)),
|
||||
XAIProfile(client_id="client", scope="openid"),
|
||||
)
|
||||
with pytest.raises(RuntimeError) as error:
|
||||
__import__("asyncio").run(
|
||||
protocol.poll_token(
|
||||
endpoint="https://auth.x.ai/oauth2/token",
|
||||
flow=type("Flow", (), {"device_code": "device", "interval": 0})(),
|
||||
timeout=1,
|
||||
)
|
||||
)
|
||||
assert "access" not in str(error.value)
|
||||
|
||||
|
||||
def test_credential_subject_prefers_id_token_subject():
|
||||
credential = XAIProtocol._credential(
|
||||
{
|
||||
"access_token": _unsigned_jwt({"principal_id": "access-subject"}),
|
||||
"refresh_token": "refresh",
|
||||
"id_token": _unsigned_jwt(
|
||||
{"sub": "stable-subject", "email": "ignored@example.test"}
|
||||
),
|
||||
"sub": "response-subject",
|
||||
},
|
||||
"https://auth.x.ai/oauth2/token",
|
||||
)
|
||||
|
||||
assert credential.subject == "stable-subject"
|
||||
|
||||
|
||||
def test_credential_subject_falls_back_to_access_token_principal_id():
|
||||
credential = XAIProtocol._credential(
|
||||
{
|
||||
"access_token": _unsigned_jwt({"principal_id": "stable-principal"}),
|
||||
"refresh_token": "refresh",
|
||||
"id_token": "malformed",
|
||||
"sub": "response-subject",
|
||||
},
|
||||
"https://auth.x.ai/oauth2/token",
|
||||
)
|
||||
|
||||
assert credential.subject == "stable-principal"
|
||||
|
||||
|
||||
def test_credential_subject_ignores_email_and_falls_back_to_response_subject():
|
||||
credential = XAIProtocol._credential(
|
||||
{
|
||||
"access_token": "not-a-jwt",
|
||||
"refresh_token": "refresh",
|
||||
"id_token": _unsigned_jwt({"email": "ignored@example.test"}),
|
||||
"sub": "response-subject",
|
||||
},
|
||||
"https://auth.x.ai/oauth2/token",
|
||||
)
|
||||
|
||||
assert credential.subject == "response-subject"
|
||||
79
tests/test_xai_enroller_sinks.py
Normal file
79
tests/test_xai_enroller_sinks.py
Normal file
@@ -0,0 +1,79 @@
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import httpx
|
||||
|
||||
from xai_enroller.models import OAuthCredential
|
||||
from xai_enroller.sinks import CPAAuthFileSink, SinkError
|
||||
|
||||
|
||||
def credential():
|
||||
return OAuthCredential(
|
||||
access_token="access",
|
||||
refresh_token="refresh",
|
||||
id_token="id-token",
|
||||
token_type="Bearer",
|
||||
expires_in=3600,
|
||||
expires_at="2026-07-11T00:00:00Z",
|
||||
last_refresh="2026-07-11T00:00:00Z",
|
||||
subject="subject",
|
||||
token_endpoint="https://auth.x.ai/oauth2/token",
|
||||
)
|
||||
|
||||
|
||||
def test_cpa_sink_builds_pinned_document_and_safe_name():
|
||||
requests = []
|
||||
|
||||
def handler(request):
|
||||
requests.append(request)
|
||||
return httpx.Response(201)
|
||||
|
||||
sink = CPAAuthFileSink(
|
||||
"https://cpa.example",
|
||||
"management-secret",
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler)),
|
||||
name_secret=b"name-secret",
|
||||
)
|
||||
receipt = asyncio.run(sink.store(credential()))
|
||||
assert requests[0].url.path == "/v0/management/auth-files"
|
||||
assert requests[0].headers["authorization"] == "Bearer management-secret"
|
||||
assert requests[0].url.params["name"].startswith("xai-")
|
||||
assert requests[0].url.params["name"].endswith(".json")
|
||||
assert json.loads(requests[0].content) == {
|
||||
"type": "xai",
|
||||
"access_token": "access",
|
||||
"refresh_token": "refresh",
|
||||
"id_token": "id-token",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600,
|
||||
"expired": "2026-07-11T00:00:00Z",
|
||||
"last_refresh": "2026-07-11T00:00:00Z",
|
||||
"sub": "subject",
|
||||
"base_url": "https://api.x.ai/v1",
|
||||
"token_endpoint": "https://auth.x.ai/oauth2/token",
|
||||
"auth_kind": "oauth",
|
||||
}
|
||||
assert "access" not in repr(receipt)
|
||||
|
||||
|
||||
def test_cpa_sink_maps_non_2xx_without_retry_or_secret_leak():
|
||||
count = 0
|
||||
|
||||
def handler(request):
|
||||
nonlocal count
|
||||
count += 1
|
||||
return httpx.Response(500, text="access")
|
||||
|
||||
sink = CPAAuthFileSink(
|
||||
"https://cpa.example",
|
||||
"management-secret",
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler)),
|
||||
name_secret=b"name-secret",
|
||||
)
|
||||
try:
|
||||
asyncio.run(sink.store(credential()))
|
||||
except SinkError as error:
|
||||
assert "access" not in str(error)
|
||||
else:
|
||||
raise AssertionError("expected SinkError")
|
||||
assert count == 1
|
||||
240
tests/test_xai_remote_snapshot.py
Normal file
240
tests/test_xai_remote_snapshot.py
Normal file
@@ -0,0 +1,240 @@
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
|
||||
from xai_enroller.models import SourceRecord
|
||||
from xai_enroller.remote_stream import (
|
||||
DiskSnapshotSource,
|
||||
LocalSnapshotSynchronizer,
|
||||
SSHSnapshotSynchronizer,
|
||||
)
|
||||
|
||||
|
||||
def _document(source_id):
|
||||
return json.dumps(
|
||||
{
|
||||
"email": source_id,
|
||||
"cookies": [
|
||||
{
|
||||
"name": "sso",
|
||||
"value": "opaque",
|
||||
"domain": "accounts.x.ai",
|
||||
"path": "/",
|
||||
}
|
||||
],
|
||||
},
|
||||
separators=(",", ":"),
|
||||
).encode()
|
||||
|
||||
|
||||
class FakeStream:
|
||||
def __init__(self, lines):
|
||||
self._lines = iter(lines)
|
||||
|
||||
async def readline(self):
|
||||
return next(self._lines, b"")
|
||||
|
||||
async def read(self, _size=-1):
|
||||
return b""
|
||||
|
||||
|
||||
class FakeProcess:
|
||||
def __init__(self, lines, returncode=0):
|
||||
self.stdout = FakeStream(lines)
|
||||
self.stderr = FakeStream([])
|
||||
self.returncode = None
|
||||
self._final_returncode = returncode
|
||||
|
||||
async def wait(self):
|
||||
self.returncode = self._final_returncode
|
||||
return self.returncode
|
||||
|
||||
def terminate(self):
|
||||
self.returncode = -15
|
||||
|
||||
def kill(self):
|
||||
self.returncode = -9
|
||||
|
||||
|
||||
def test_snapshot_sync_atomically_replaces_only_valid_complete_export(tmp_path):
|
||||
async def scenario():
|
||||
destination = tmp_path / "source-snapshot.jsonl"
|
||||
destination.write_bytes(_document("old") + b"\n")
|
||||
|
||||
invalid = FakeProcess([b"not-json\n"])
|
||||
|
||||
async def invalid_factory(*_args, **_kwargs):
|
||||
return invalid
|
||||
|
||||
synchronizer = SSHSnapshotSynchronizer(
|
||||
"host",
|
||||
destination,
|
||||
process_factory=invalid_factory,
|
||||
fingerprint=lambda source_id: f"key:{source_id}",
|
||||
)
|
||||
assert not await synchronizer.sync_once()
|
||||
assert destination.read_bytes() == _document("old") + b"\n"
|
||||
|
||||
rejected = FakeProcess([_document("rejected") + b"\n"], returncode=1)
|
||||
|
||||
async def rejected_factory(*_args, **_kwargs):
|
||||
return rejected
|
||||
|
||||
synchronizer.process_factory = rejected_factory
|
||||
assert not await synchronizer.sync_once()
|
||||
assert destination.read_bytes() == _document("old") + b"\n"
|
||||
|
||||
empty = FakeProcess([])
|
||||
|
||||
async def empty_factory(*_args, **_kwargs):
|
||||
return empty
|
||||
|
||||
synchronizer.process_factory = empty_factory
|
||||
assert not await synchronizer.sync_once()
|
||||
assert destination.read_bytes() == _document("old") + b"\n"
|
||||
|
||||
valid = FakeProcess([_document("new-a") + b"\n", _document("new-b") + b"\n"])
|
||||
|
||||
async def valid_factory(*_args, **_kwargs):
|
||||
return valid
|
||||
|
||||
synchronizer.process_factory = valid_factory
|
||||
assert await synchronizer.sync_once()
|
||||
assert synchronizer.snapshot_fingerprints == frozenset(
|
||||
{"key:new-a", "key:new-b"}
|
||||
)
|
||||
assert destination.read_bytes() == _document("new-a") + b"\n" + _document("new-b") + b"\n"
|
||||
assert destination.stat().st_mode & 0o777 == 0o600
|
||||
assert not list(tmp_path.glob(".source-snapshot.jsonl.*.tmp"))
|
||||
|
||||
synchronizer.process_factory = invalid_factory
|
||||
assert not await synchronizer.sync_once()
|
||||
assert synchronizer.snapshot_fingerprints == frozenset(
|
||||
{"key:new-a", "key:new-b"}
|
||||
)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_local_snapshot_accepts_empty_input_and_anchors_exporter_code(tmp_path):
|
||||
async def scenario():
|
||||
registration_root = tmp_path / "registration"
|
||||
destination = tmp_path / "auth" / "source-snapshot.jsonl"
|
||||
captured = []
|
||||
process = FakeProcess([])
|
||||
|
||||
async def factory(*args, **_kwargs):
|
||||
captured.append(args)
|
||||
return process
|
||||
|
||||
synchronizer = LocalSnapshotSynchronizer(
|
||||
registration_root,
|
||||
destination,
|
||||
process_factory=factory,
|
||||
)
|
||||
|
||||
assert await synchronizer.sync_once()
|
||||
assert destination.read_bytes() == b""
|
||||
assert destination.stat().st_mode & 0o777 == 0o600
|
||||
assert synchronizer.snapshot_fingerprints == frozenset()
|
||||
assert captured[0][0]
|
||||
assert captured[0][1] != str(registration_root / "scripts" / "export_registered_sessions.py")
|
||||
assert captured[0][-2:] == (
|
||||
str(registration_root / "keys" / "auth-sessions.jsonl"),
|
||||
str(registration_root / "keys" / "accounts.txt"),
|
||||
)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_local_snapshot_rejects_candidate_when_either_input_changes(tmp_path):
|
||||
async def scenario():
|
||||
registration_root = tmp_path / "registration"
|
||||
keys = registration_root / "keys"
|
||||
keys.mkdir(parents=True)
|
||||
sessions = keys / "auth-sessions.jsonl"
|
||||
accounts = keys / "accounts.txt"
|
||||
sessions.write_bytes(_document("old") + b"\n")
|
||||
accounts.write_text("", encoding="utf-8")
|
||||
destination = tmp_path / "auth" / "source-snapshot.jsonl"
|
||||
destination.parent.mkdir()
|
||||
destination.write_bytes(_document("stable") + b"\n")
|
||||
|
||||
process = FakeProcess([_document("candidate") + b"\n"])
|
||||
|
||||
async def factory(*_args, **_kwargs):
|
||||
accounts.write_text("new@example.test:p:sso\n", encoding="utf-8")
|
||||
return process
|
||||
|
||||
synchronizer = LocalSnapshotSynchronizer(
|
||||
registration_root,
|
||||
destination,
|
||||
process_factory=factory,
|
||||
)
|
||||
|
||||
assert not await synchronizer.sync_once()
|
||||
assert destination.read_bytes() == _document("stable") + b"\n"
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_snapshot_consumer_finishes_open_generation_before_replacement(tmp_path):
|
||||
async def scenario():
|
||||
destination = tmp_path / "source-snapshot.jsonl"
|
||||
destination.write_bytes(_document("old-a") + b"\n" + _document("old-b") + b"\n")
|
||||
source = DiskSnapshotSource(destination, poll_seconds=0.01)
|
||||
records = source.records()
|
||||
first = await anext(records)
|
||||
assert isinstance(first, SourceRecord)
|
||||
assert first.source_id == "old-a"
|
||||
|
||||
replacement = tmp_path / "replacement"
|
||||
replacement.write_bytes(_document("new-a") + b"\n")
|
||||
os.replace(replacement, destination)
|
||||
|
||||
assert (await anext(records)).source_id == "old-b"
|
||||
assert (await asyncio.wait_for(anext(records), 1)).source_id == "new-a"
|
||||
await source.close()
|
||||
await records.aclose()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_snapshot_source_reports_only_aggregate_new_records(tmp_path):
|
||||
class Synchronizer:
|
||||
def __init__(self):
|
||||
self.snapshot_fingerprints = None
|
||||
self.calls = 0
|
||||
self.source = None
|
||||
|
||||
async def sync_once(self):
|
||||
self.calls += 1
|
||||
if self.calls == 1:
|
||||
self.snapshot_fingerprints = frozenset({"a", "b"})
|
||||
else:
|
||||
self.snapshot_fingerprints = frozenset({"a", "b", "c"})
|
||||
self.source._closed = True
|
||||
return True
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
async def scenario():
|
||||
synchronizer = Synchronizer()
|
||||
events = []
|
||||
source = DiskSnapshotSource(
|
||||
tmp_path / "source-snapshot.jsonl",
|
||||
synchronizer=synchronizer,
|
||||
sync_seconds=0.001,
|
||||
event_callback=lambda kind, data: events.append((kind, data)),
|
||||
)
|
||||
synchronizer.source = source
|
||||
|
||||
await source._sync_loop()
|
||||
|
||||
assert events == [
|
||||
("source_connected", {}),
|
||||
("source_updated", {"new": 1, "total": 3}),
|
||||
]
|
||||
|
||||
asyncio.run(scenario())
|
||||
Reference in New Issue
Block a user