mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-05 04:13:25 +00:00
Alex
574f96341e
The backend import package is now docsgpt, the name it will carry on PyPI; application was far too generic to install into anyone's site-packages. git mv plus a mechanical rewrite of every import, dotted string and path reference: 734 Python files, the compose files, Dockerfile, workflows, docs, setup scripts, devcontainer, k8s manifests, vscode config, pytest and coverage config, .gitignore. Behaviour is unchanged. Kept for one release: - A top-level application package whose meta-path finder resolves application.x.y to the already-imported docsgpt.x.y object, so old imports and entry points (celery -A application.app.celery, uvicorn application.asgi:asgi_app) keep working with a FutureWarning. - Celery registers every application.* task name as an alias of its docsgpt.* task on start-up, so messages queued by the previous release still run. The redbeat key prefix moves to redbeat:docsgpt:v2: so schedule entries the previous release wrote are left unread instead of firing twice. The backend image builds from the repository root (docker build -f docsgpt/Dockerfile .) so it can ship the alias package; a root .dockerignore allow-lists docsgpt/ and application/ and keeps caches, local data, .env files, the sample index files and the Dockerfile out. Compose and the image workflows point at the new context.
229 lines
7.4 KiB
Python
229 lines
7.4 KiB
Python
"""Tests for poll-issued session-ticket enforcement on the SSE upgrade.
|
|
|
|
Two layers:
|
|
|
|
* Broker unit tests for ``claim_ticket`` / ``validate_ticket`` (issue, match,
|
|
mismatch, eviction, absence).
|
|
* Route tests proving the real CLI loop still works: ``/poll`` issues a
|
|
ticket and ``session_events`` accepts *that* ticket, while a mismatched
|
|
ticket is rejected with ``410`` before any stream opens.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from flask import Flask
|
|
|
|
from docsgpt.api.devices import auth as auth_module
|
|
from docsgpt.api.devices import session as session_module
|
|
from docsgpt.devices.broker import DeviceBroker
|
|
|
|
from .conftest import FakeRedis
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Broker unit tests
|
|
# ---------------------------------------------------------------------------
|
|
def _queue_work(broker: DeviceBroker, device_id: str = "dev_x") -> None:
|
|
# Queuing a command (no draining session) is what makes claim_ticket
|
|
# hand out a ticket: the device's command list is non-empty.
|
|
broker.dispatch_invocation(
|
|
device_id, "user_x", {"invocation_id": "inv_1", "action": "run_command"}
|
|
)
|
|
|
|
|
|
def test_claim_ticket_none_when_no_work(broker_env):
|
|
broker, _fake = broker_env
|
|
assert broker.claim_ticket("dev_x", 30) is None
|
|
|
|
|
|
def test_validate_accepts_issued_ticket(broker_env):
|
|
broker, _fake = broker_env
|
|
_queue_work(broker)
|
|
ticket = broker.claim_ticket("dev_x", 30)
|
|
assert ticket is not None
|
|
assert broker.validate_ticket("dev_x", ticket) is True
|
|
|
|
|
|
def test_validate_rejects_wrong_ticket(broker_env):
|
|
broker, _fake = broker_env
|
|
_queue_work(broker)
|
|
broker.claim_ticket("dev_x", 30)
|
|
assert broker.validate_ticket("dev_x", "st_not_the_one") is False
|
|
|
|
|
|
def test_validate_rejects_when_no_ticket_issued(broker_env):
|
|
broker, _fake = broker_env
|
|
assert broker.validate_ticket("dev_x", "st_anything") is False
|
|
|
|
|
|
def test_validate_rejects_empty_session_id(broker_env):
|
|
broker, _fake = broker_env
|
|
_queue_work(broker)
|
|
broker.claim_ticket("dev_x", 30)
|
|
assert broker.validate_ticket("dev_x", "") is False
|
|
|
|
|
|
def test_validate_rejects_after_ticket_evicted(broker_env):
|
|
# Redis enforces the TTL; once the ticket key is gone (expired/evicted),
|
|
# validate can't resurrect it.
|
|
broker, fake = broker_env
|
|
_queue_work(broker)
|
|
ticket = broker.claim_ticket("dev_x", 30)
|
|
fake.delete("dev:ticket:dev_x") # simulate TTL expiry
|
|
assert broker.validate_ticket("dev_x", ticket) is False
|
|
|
|
|
|
def test_claim_ticket_reused_while_unexpired(broker_env):
|
|
broker, _fake = broker_env
|
|
_queue_work(broker)
|
|
first = broker.claim_ticket("dev_x", 30)
|
|
second = broker.claim_ticket("dev_x", 30)
|
|
assert first == second
|
|
|
|
|
|
def test_register_session_consumes_issued_ticket_as_session_id(broker_env):
|
|
# The issued ticket becomes the session_id, so the URL the CLI opens
|
|
# (= the ticket) matches the live session.
|
|
broker, fake = broker_env
|
|
_queue_work(broker)
|
|
ticket = broker.claim_ticket("dev_x", 30)
|
|
sess = broker.register_session("dev_x", "user_x")
|
|
assert sess.session_id == ticket
|
|
# Ticket is consumed on registration.
|
|
assert fake.get("dev:ticket:dev_x") is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Route tests (poll -> SSE upgrade)
|
|
# ---------------------------------------------------------------------------
|
|
@pytest.fixture
|
|
def app():
|
|
return Flask(__name__)
|
|
|
|
|
|
def _device_row() -> dict:
|
|
return {
|
|
"id": "dev_route",
|
|
"user_id": "user_route",
|
|
"name": "laptop",
|
|
"status": "active",
|
|
"approval_mode": "ask",
|
|
"machine_pubkey_fingerprint": "fp",
|
|
"token_hash": "tokhash",
|
|
}
|
|
|
|
|
|
class _Repo:
|
|
def __init__(self, _conn):
|
|
pass
|
|
|
|
def find_by_token_hash(self, _token_hash):
|
|
return _device_row()
|
|
|
|
def touch_last_seen(self, _device_id):
|
|
pass
|
|
|
|
|
|
class _Ctx:
|
|
def __enter__(self):
|
|
return None
|
|
|
|
def __exit__(self, *a):
|
|
return False
|
|
|
|
|
|
def _patched_auth():
|
|
return [
|
|
patch.object(auth_module, "DevicesRepository", _Repo),
|
|
patch.object(auth_module, "db_readonly", _Ctx),
|
|
patch.object(auth_module, "db_session", _Ctx),
|
|
]
|
|
|
|
|
|
def _call(app: Flask, view, path: str, *args):
|
|
with app.test_request_context(
|
|
path, method="GET", headers={"Authorization": "Bearer tok_good"}
|
|
):
|
|
ctxs = _patched_auth()
|
|
for c in ctxs:
|
|
c.start()
|
|
try:
|
|
return view(*args)
|
|
finally:
|
|
for c in ctxs:
|
|
c.stop()
|
|
|
|
|
|
def test_poll_to_sse_upgrade_with_issued_ticket_works(app):
|
|
"""The legitimate loop: /poll issues a ticket, session_events accepts it."""
|
|
fake = FakeRedis()
|
|
broker = DeviceBroker()
|
|
with patch("docsgpt.devices.broker.get_redis_instance", return_value=fake):
|
|
# Queue work so /poll returns a ticket rather than 202.
|
|
broker.dispatch_invocation(
|
|
"dev_route", "user_route",
|
|
{"invocation_id": "inv_a", "action": "run_command"},
|
|
)
|
|
with patch.object(session_module, "get_broker", return_value=broker):
|
|
poll_resp = _call(app, session_module.poll, "/api/devices/poll")
|
|
assert poll_resp.status_code == 200
|
|
payload = poll_resp.get_json()
|
|
ticket = payload["session_ticket"]
|
|
assert payload["session_url"] == f"/api/devices/sessions/{ticket}/events"
|
|
assert payload["expires_in"] == 30
|
|
|
|
# CLI opens the exact session_url it was handed.
|
|
sse_resp = _call(
|
|
app,
|
|
session_module.session_events,
|
|
f"/api/devices/sessions/{ticket}/events",
|
|
ticket,
|
|
)
|
|
try:
|
|
assert sse_resp.status_code == 200
|
|
assert sse_resp.mimetype == "text/event-stream"
|
|
finally:
|
|
sse_resp.close()
|
|
|
|
|
|
def test_session_events_rejects_mismatched_ticket(app):
|
|
"""A fabricated/mismatched session_id is 410 Gone, no stream opened."""
|
|
fake = FakeRedis()
|
|
broker = DeviceBroker()
|
|
with patch("docsgpt.devices.broker.get_redis_instance", return_value=fake):
|
|
broker.dispatch_invocation(
|
|
"dev_route", "user_route",
|
|
{"invocation_id": "inv_b", "action": "run_command"},
|
|
)
|
|
with patch.object(session_module, "get_broker", return_value=broker):
|
|
# Poll issues the real ticket...
|
|
_call(app, session_module.poll, "/api/devices/poll")
|
|
# ...but the client opens a different one.
|
|
resp = _call(
|
|
app,
|
|
session_module.session_events,
|
|
"/api/devices/sessions/st_bogus/events",
|
|
"st_bogus",
|
|
)
|
|
assert resp.status_code == 410
|
|
assert resp.get_json()["error"] == "session_ticket_invalid"
|
|
|
|
|
|
def test_session_events_rejects_when_never_polled(app):
|
|
"""Opening the SSE stream without a prior poll is rejected (410)."""
|
|
fake = FakeRedis()
|
|
broker = DeviceBroker()
|
|
with patch("docsgpt.devices.broker.get_redis_instance", return_value=fake):
|
|
with patch.object(session_module, "get_broker", return_value=broker):
|
|
resp = _call(
|
|
app,
|
|
session_module.session_events,
|
|
"/api/devices/sessions/st_anything/events",
|
|
"st_anything",
|
|
)
|
|
assert resp.status_code == 410
|
|
assert resp.get_json()["error"] == "session_ticket_invalid"
|