mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 18:18:08 +00:00
Redundant test
This commit is contained in:
1 parent
6c63c91643
commit
f878a0f902
5 files changed
+89
-32
No files matched your search
@@ -29,6 +29,9 @@ from application.api.user.team_sharing import (
|
||||
from application.storage.db.repositories.agent_folders import AgentFoldersRepository
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.conversations import ConversationsRepository
|
||||
from application.storage.db.repositories.shared_conversations import (
|
||||
SharedConversationsRepository,
|
||||
)
|
||||
from application.storage.db.repositories.stack_logs import StackLogsRepository
|
||||
from application.storage.db.repositories.token_usage import TokenUsageRepository
|
||||
from application.storage.db.repositories.users import UsersRepository
|
||||
@@ -1360,13 +1363,16 @@ class RegenerateAgentKey(Resource):
|
||||
500,
|
||||
)
|
||||
|
||||
# Re-point key-string history to the new key so per-agent
|
||||
# analytics survive rotation. token_usage/stack_logs (0026)
|
||||
# carry agent_id, but rows can still be api_key-only:
|
||||
# Re-point every api_key reference to the new key so rotation
|
||||
# doesn't break history or live links. token_usage/stack_logs
|
||||
# (0026) carry agent_id, but rows can still be api_key-only:
|
||||
# conversations set api_key without an agent_id unless the caller
|
||||
# passes one, stack_logs rows may have a NULL agent_id, and the
|
||||
# 24h rate-limit window is keyed by api_key only (token_usage).
|
||||
# All three rewrites run in the same transaction as the swap.
|
||||
# shared_conversations is functional, not just analytics: the
|
||||
# public share endpoint hands its stored api_key to the widget,
|
||||
# so a stale key would break promptable shared links.
|
||||
# All rewrites run in the same transaction as the key swap.
|
||||
StackLogsRepository(conn).reassign_api_key(
|
||||
old_key=old_key, new_key=new_key
|
||||
)
|
||||
@@ -1376,6 +1382,9 @@ class RegenerateAgentKey(Resource):
|
||||
ConversationsRepository(conn).reassign_api_key(
|
||||
old_key=old_key, new_key=new_key
|
||||
)
|
||||
SharedConversationsRepository(conn).reassign_api_key(
|
||||
old_key=old_key, new_key=new_key
|
||||
)
|
||||
except Exception as err:
|
||||
current_app.logger.error(
|
||||
f"Error regenerating agent key: {err}", exc_info=True
|
||||
|
||||
@@ -211,3 +211,31 @@ class SharedConversationsRepository:
|
||||
{"conv_id": conversation_id},
|
||||
)
|
||||
return [row_to_dict(r) for r in result.fetchall()]
|
||||
|
||||
def reassign_api_key(self, *, old_key: str, new_key: str) -> int:
|
||||
"""Re-point promptable share rows from ``old_key`` to ``new_key``.
|
||||
|
||||
A *promptable* shared-conversation link stores the backing agent's key
|
||||
in ``api_key`` and the public share endpoint returns it so the shared
|
||||
widget can keep prompting. If that agent's key is rotated without this
|
||||
rewrite, the endpoint would hand out an invalidated key and break every
|
||||
existing shared link. Rewriting keeps those links working with the new
|
||||
key.
|
||||
|
||||
Matched by ``api_key`` only — ``agents.key`` is globally unique, so a
|
||||
key maps to exactly one agent. The uniform rewrite cannot violate the
|
||||
``(conversation_id, user_id, is_promptable, first_n_queries,
|
||||
COALESCE(api_key, ''))`` partial unique index: rows already differ on
|
||||
the non-key columns, and ``new_key`` is freshly generated so no row
|
||||
already holds it. Returns the number of rows updated.
|
||||
"""
|
||||
if not old_key or not new_key:
|
||||
return 0
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"UPDATE shared_conversations SET api_key = :new_key "
|
||||
"WHERE api_key = :old_key"
|
||||
),
|
||||
{"old_key": old_key, "new_key": new_key},
|
||||
)
|
||||
return result.rowcount
|
||||
@@ -1343,16 +1343,20 @@ class TestRegenerateAgentKey:
|
||||
assert repo.find_by_key(new_key)["id"] == agent["id"]
|
||||
|
||||
def test_migrates_historical_logs_and_usage(self, app, pg_conn):
|
||||
"""stack_logs + token_usage + conversations rows follow the key so
|
||||
analytics and the rate-limit window are not orphaned. The stack_logs row
|
||||
is stamped with a *different* user_id (a caller, not the owner) to prove
|
||||
the migration is not owner-scoped, and the conversation is created with
|
||||
api_key set but agent_id NULL to prove the api_key-only path is covered.
|
||||
"""stack_logs + token_usage + conversations + shared_conversations rows
|
||||
follow the key so analytics, the rate-limit window, and promptable shared
|
||||
links are not orphaned/broken. The stack_logs row is stamped with a
|
||||
*different* user_id (a caller, not the owner) to prove the migration is
|
||||
not owner-scoped, and the conversation is created with api_key set but
|
||||
agent_id NULL to prove the api_key-only path is covered.
|
||||
"""
|
||||
from application.api.user.agents.routes import RegenerateAgentKey
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
from application.storage.db.repositories.shared_conversations import (
|
||||
SharedConversationsRepository,
|
||||
)
|
||||
from application.storage.db.repositories.stack_logs import (
|
||||
StackLogsRepository,
|
||||
)
|
||||
@@ -1375,9 +1379,14 @@ class TestRegenerateAgentKey:
|
||||
api_key=old_key, prompt_tokens=7, generated_tokens=3,
|
||||
)
|
||||
# api_key set, agent_id intentionally omitted (NULL) — the orphan case.
|
||||
ConversationsRepository(pg_conn).create(
|
||||
conv = ConversationsRepository(pg_conn).create(
|
||||
user, "conv", api_key=old_key,
|
||||
)
|
||||
# A promptable shared link stores the agent key; rotation must not
|
||||
# leave it pointing at an invalidated key.
|
||||
SharedConversationsRepository(pg_conn).create(
|
||||
conv["id"], user, is_promptable=True, api_key=old_key,
|
||||
)
|
||||
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/regenerate_agent_key/{agent['id']}", method="POST"
|
||||
@@ -1402,6 +1411,8 @@ class TestRegenerateAgentKey:
|
||||
assert _count("token_usage", new_key) == 1
|
||||
assert _count("conversations", old_key) == 0
|
||||
assert _count("conversations", new_key) == 1
|
||||
assert _count("shared_conversations", old_key) == 0
|
||||
assert _count("shared_conversations", new_key) == 1
|
||||
|
||||
def test_db_error_returns_500(self, app):
|
||||
from application.api.user.agents.routes import RegenerateAgentKey
|
||||
|
||||
@@ -115,3 +115,34 @@ class TestFindByUuidShapeGate:
|
||||
share = repo.create(conv["id"], "user-1", first_n_queries=1)
|
||||
found = repo.find_by_uuid(str(share["uuid"]))
|
||||
assert found is not None
|
||||
|
||||
|
||||
class TestReassignApiKey:
|
||||
def test_rewrites_promptable_share_keys(self, pg_conn):
|
||||
# A promptable share stores the backing agent's key; on rotation it must
|
||||
# follow the key so the public share endpoint keeps returning a valid one.
|
||||
from sqlalchemy import text
|
||||
|
||||
repo = _repo(pg_conn)
|
||||
c1, c2, c3 = _conv(pg_conn), _conv(pg_conn), _conv(pg_conn)
|
||||
repo.create(c1["id"], "u", is_promptable=True, api_key="sc-old")
|
||||
repo.create(c2["id"], "u", is_promptable=True, api_key="sc-old")
|
||||
repo.create(c3["id"], "u", is_promptable=True, api_key="sc-other")
|
||||
|
||||
moved = repo.reassign_api_key(old_key="sc-old", new_key="sc-new")
|
||||
assert moved == 2
|
||||
|
||||
def _count(key):
|
||||
return pg_conn.execute(
|
||||
text("SELECT COUNT(*) FROM shared_conversations WHERE api_key = :k"),
|
||||
{"k": key},
|
||||
).scalar()
|
||||
|
||||
assert _count("sc-old") == 0
|
||||
assert _count("sc-new") == 2
|
||||
assert _count("sc-other") == 1
|
||||
|
||||
def test_noop_on_blank_keys(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
assert repo.reassign_api_key(old_key="", new_key="x") == 0
|
||||
assert repo.reassign_api_key(old_key="x", new_key="") == 0
|
||||
@@ -199,28 +199,6 @@ class TestInsert:
|
||||
assert dict(row._mapping)["stacks"][0]["data"]["result"] == "ok"
|
||||
|
||||
|
||||
class TestAgentId:
|
||||
def test_valid_agent_id_stored(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
agent_id = "11111111-2222-3333-4444-555555555555"
|
||||
repo.insert(activity_id="ai-valid", agent_id=agent_id)
|
||||
row = pg_conn.execute(
|
||||
text("SELECT agent_id FROM stack_logs WHERE activity_id = 'ai-valid'")
|
||||
).fetchone()
|
||||
assert str(dict(row._mapping)["agent_id"]) == agent_id
|
||||
|
||||
def test_invalid_agent_id_coerced_to_null(self, pg_conn):
|
||||
# A 24-hex legacy Mongo ObjectId is not a valid UUID; it must be
|
||||
# coerced to NULL rather than reaching CAST(... AS uuid) and killing
|
||||
# the whole activity-log write.
|
||||
repo = _repo(pg_conn)
|
||||
repo.insert(activity_id="ai-legacy", agent_id="507f1f77bcf86cd799439011")
|
||||
row = pg_conn.execute(
|
||||
text("SELECT agent_id FROM stack_logs WHERE activity_id = 'ai-legacy'")
|
||||
).fetchone()
|
||||
assert dict(row._mapping)["agent_id"] is None
|
||||
|
||||
|
||||
class TestReassignApiKey:
|
||||
def test_rewrites_matching_rows_regardless_of_user(self, pg_conn):
|
||||
# On key rotation the rewrite is intentionally NOT scoped by user_id:
|
||||
|
||||
Reference in new issue
Block a user