d10009d639
Open-source Grok free registration CLI, xai_enroller auth pipeline, local auth service, tests and docs.
299 lines
8.5 KiB
Python
299 lines
8.5 KiB
Python
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())
|