mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 14:15:52 +00:00
fix: tests
This commit is contained in:
Alex committed
1 parent
81b6ee5daa
commit
ee430aff1e
5 files changed
-1328
No files matched your search
@@ -1,164 +0,0 @@
|
||||
"""Integration tests for ``_backfill_connector_sessions``.
|
||||
|
||||
Mongo routinely contains multiple rows for the same
|
||||
``(user_id, server_url, provider)`` triple — each OAuth-button click
|
||||
inserts a pending row and only the last one is authorized. The Postgres
|
||||
schema has a unique index on that triple, so the backfill does a Python-
|
||||
side dedup *before* the insert, preferring authorized rows over pending
|
||||
ones and newer ``created_at`` over older.
|
||||
|
||||
This test seeds exactly that shape and asserts the dedup keeps the
|
||||
authorized row.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import mongomock
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[3]))
|
||||
|
||||
from scripts.db.backfill import _backfill_connector_sessions # noqa: E402
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mongo_db() -> Any:
|
||||
client = mongomock.MongoClient()
|
||||
return client["docsgpt_test"]
|
||||
|
||||
|
||||
class TestBackfillConnectorSessions:
|
||||
def test_dedups_authorized_over_pending(self, pg_conn, mongo_db):
|
||||
# Two rows, same (user, server_url, provider) triple. The pending
|
||||
# row was inserted first; the authorized row was inserted later
|
||||
# when the OAuth redirect completed. Dedup should keep the
|
||||
# authorized one (has ``token_info``).
|
||||
mongo_db["connector_sessions"].insert_many(
|
||||
[
|
||||
{
|
||||
"_id": "777777777777777777777771",
|
||||
"user_id": "alice",
|
||||
"provider": "google_drive",
|
||||
"server_url": "https://drive.google.com",
|
||||
"status": "pending",
|
||||
"created_at": datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||||
},
|
||||
{
|
||||
"_id": "777777777777777777777772",
|
||||
"user_id": "alice",
|
||||
"provider": "google_drive",
|
||||
"server_url": "https://drive.google.com",
|
||||
"status": "authorized",
|
||||
"token_info": {"access_token": "xyz"},
|
||||
"created_at": datetime(2026, 1, 2, tzinfo=timezone.utc),
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
stats = _backfill_connector_sessions(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
assert stats["seen"] == 2
|
||||
assert stats["written"] == 1
|
||||
assert stats["skipped"] == 1
|
||||
|
||||
rows = pg_conn.execute(
|
||||
text(
|
||||
"SELECT status, legacy_mongo_id, token_info FROM "
|
||||
"connector_sessions WHERE user_id = 'alice'"
|
||||
)
|
||||
).fetchall()
|
||||
assert len(rows) == 1
|
||||
assert rows[0]._mapping["status"] == "authorized"
|
||||
assert rows[0]._mapping["legacy_mongo_id"] == "777777777777777777777772"
|
||||
assert rows[0]._mapping["token_info"] == {"access_token": "xyz"}
|
||||
|
||||
def test_dedups_newer_created_at_when_both_pending(self, pg_conn, mongo_db):
|
||||
# Both rows are pending. Newer created_at wins.
|
||||
mongo_db["connector_sessions"].insert_many(
|
||||
[
|
||||
{
|
||||
"_id": "777777777777777777777771",
|
||||
"user_id": "alice",
|
||||
"provider": "github",
|
||||
"server_url": "",
|
||||
"status": "pending",
|
||||
"created_at": datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||||
},
|
||||
{
|
||||
"_id": "777777777777777777777772",
|
||||
"user_id": "alice",
|
||||
"provider": "github",
|
||||
"server_url": "",
|
||||
"status": "pending",
|
||||
"created_at": datetime(2026, 2, 1, tzinfo=timezone.utc),
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
_backfill_connector_sessions(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
rows = pg_conn.execute(
|
||||
text(
|
||||
"SELECT legacy_mongo_id FROM connector_sessions "
|
||||
"WHERE user_id = 'alice' AND provider = 'github'"
|
||||
)
|
||||
).fetchall()
|
||||
assert len(rows) == 1
|
||||
assert rows[0]._mapping["legacy_mongo_id"] == "777777777777777777777772"
|
||||
|
||||
def test_rerun_does_not_duplicate(self, pg_conn, mongo_db):
|
||||
mongo_db["connector_sessions"].insert_one(
|
||||
{
|
||||
"_id": "777777777777777777777771",
|
||||
"user_id": "alice",
|
||||
"provider": "google_drive",
|
||||
"server_url": "https://drive.google.com",
|
||||
"status": "authorized",
|
||||
"token_info": {"access_token": "xyz"},
|
||||
"created_at": datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||||
}
|
||||
)
|
||||
_backfill_connector_sessions(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
_backfill_connector_sessions(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
count = pg_conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM connector_sessions "
|
||||
"WHERE legacy_mongo_id = '777777777777777777777771'"
|
||||
)
|
||||
).scalar()
|
||||
assert count == 1
|
||||
|
||||
def test_skips_rows_without_user_or_provider(self, pg_conn, mongo_db):
|
||||
mongo_db["connector_sessions"].insert_many(
|
||||
[
|
||||
{
|
||||
"_id": "777777777777777777777771",
|
||||
"provider": "github", # no user
|
||||
},
|
||||
{
|
||||
"_id": "777777777777777777777772",
|
||||
"user_id": "alice", # no provider
|
||||
},
|
||||
]
|
||||
)
|
||||
stats = _backfill_connector_sessions(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
assert stats["seen"] == 2
|
||||
assert stats["written"] == 0
|
||||
assert stats["skipped"] == 2
|
||||
@@ -1,292 +0,0 @@
|
||||
"""Integration tests for ``_backfill_conversations`` +
|
||||
``_backfill_conversation_messages`` flattening and attachment FK
|
||||
resolution.
|
||||
|
||||
The conversations backfill does two things the other backfillers don't:
|
||||
1. It flattens Mongo's nested ``queries[]`` array into rows in a child
|
||||
table (``conversation_messages``), with ``position`` = array index.
|
||||
2. It resolves attachment refs (strings pointing to Mongo ObjectIds) to
|
||||
Postgres UUIDs via a ``legacy_mongo_id`` lookup. Unresolvable refs are
|
||||
dropped rather than crashing the whole batch.
|
||||
|
||||
These are the two bits we assert below, alongside the standard
|
||||
happy-shape / idempotency / convergence checks.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import mongomock
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[3]))
|
||||
|
||||
from scripts.db.backfill import ( # noqa: E402
|
||||
_backfill_attachments,
|
||||
_backfill_conversations,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mongo_db() -> Any:
|
||||
client = mongomock.MongoClient()
|
||||
return client["docsgpt_test"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# attachments — prerequisite for conversations (DBRef→UUID via legacy map)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBackfillAttachments:
|
||||
def test_attachments_happy_shape_preserves_mime_and_size(
|
||||
self, pg_conn, mongo_db
|
||||
):
|
||||
mongo_db["attachments"].insert_one(
|
||||
{
|
||||
"_id": "aaaaaaaaaaaaaaaaaaaaaaaa",
|
||||
"user": "alice",
|
||||
"filename": "report.pdf",
|
||||
# Worker writes the blob path as ``path``; the PG column is
|
||||
# ``upload_path``. If the backfill ever reads the wrong key,
|
||||
# this test will fail with an empty string.
|
||||
"path": "uploads/alice/report.pdf",
|
||||
"mime_type": "application/pdf",
|
||||
"size": 12345,
|
||||
"content": "extracted text",
|
||||
"token_count": 42,
|
||||
}
|
||||
)
|
||||
|
||||
_backfill_attachments(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT filename, upload_path, mime_type, size, "
|
||||
"content, token_count, legacy_mongo_id "
|
||||
"FROM attachments WHERE user_id = 'alice'"
|
||||
)
|
||||
).one()._mapping
|
||||
assert row["filename"] == "report.pdf"
|
||||
assert row["upload_path"] == "uploads/alice/report.pdf"
|
||||
assert row["mime_type"] == "application/pdf"
|
||||
assert row["size"] == 12345
|
||||
assert row["content"] == "extracted text"
|
||||
assert row["token_count"] == 42
|
||||
assert row["legacy_mongo_id"] == "aaaaaaaaaaaaaaaaaaaaaaaa"
|
||||
|
||||
def test_attachments_rerun_does_not_duplicate(self, pg_conn, mongo_db):
|
||||
mongo_db["attachments"].insert_one(
|
||||
{
|
||||
"_id": "aaaaaaaaaaaaaaaaaaaaaaaa",
|
||||
"user": "alice",
|
||||
"filename": "r.pdf",
|
||||
"path": "u/alice/r.pdf",
|
||||
}
|
||||
)
|
||||
_backfill_attachments(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
_backfill_attachments(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
count = pg_conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM attachments "
|
||||
"WHERE legacy_mongo_id = 'aaaaaaaaaaaaaaaaaaaaaaaa'"
|
||||
)
|
||||
).scalar()
|
||||
assert count == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# conversations + conversation_messages
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _seed_attachment(mongo_db: Any, _id: str, user: str = "alice") -> None:
|
||||
mongo_db["attachments"].insert_one(
|
||||
{
|
||||
"_id": _id,
|
||||
"user": user,
|
||||
"filename": f"{_id}.txt",
|
||||
"path": f"uploads/{user}/{_id}.txt",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class TestBackfillConversations:
|
||||
def test_conversations_flattens_queries_into_messages(
|
||||
self, pg_conn, mongo_db
|
||||
):
|
||||
_seed_attachment(mongo_db, "aaaaaaaaaaaaaaaaaaaaaaa1")
|
||||
_backfill_attachments(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
mongo_db["conversations"].insert_one(
|
||||
{
|
||||
"_id": "cccccccccccccccccccccccc",
|
||||
"user": "alice",
|
||||
"name": "chat-1",
|
||||
"queries": [
|
||||
{
|
||||
"prompt": "hello",
|
||||
"response": "hi",
|
||||
"attachments": ["aaaaaaaaaaaaaaaaaaaaaaa1"],
|
||||
},
|
||||
{
|
||||
"prompt": "how are you",
|
||||
"response": "fine",
|
||||
},
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
stats = _backfill_conversations(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
assert stats["seen"] == 1
|
||||
assert stats["written"] == 1
|
||||
assert stats["messages_written"] == 2
|
||||
|
||||
conv_id = pg_conn.execute(
|
||||
text(
|
||||
"SELECT id FROM conversations "
|
||||
"WHERE legacy_mongo_id = 'cccccccccccccccccccccccc'"
|
||||
)
|
||||
).scalar()
|
||||
|
||||
rows = pg_conn.execute(
|
||||
text(
|
||||
"SELECT position, prompt, response, attachments "
|
||||
"FROM conversation_messages WHERE conversation_id = :cid "
|
||||
"ORDER BY position"
|
||||
),
|
||||
{"cid": str(conv_id)},
|
||||
).fetchall()
|
||||
assert [r._mapping["position"] for r in rows] == [0, 1]
|
||||
assert rows[0]._mapping["prompt"] == "hello"
|
||||
assert rows[1]._mapping["response"] == "fine"
|
||||
# Attachment ObjectId string was mapped to the PG UUID.
|
||||
resolved = rows[0]._mapping["attachments"]
|
||||
assert len(resolved) == 1
|
||||
pg_att_id = pg_conn.execute(
|
||||
text(
|
||||
"SELECT id FROM attachments "
|
||||
"WHERE legacy_mongo_id = 'aaaaaaaaaaaaaaaaaaaaaaa1'"
|
||||
)
|
||||
).scalar()
|
||||
assert str(resolved[0]) == str(pg_att_id)
|
||||
|
||||
def test_conversations_drops_unresolved_attachments(
|
||||
self, pg_conn, mongo_db
|
||||
):
|
||||
mongo_db["conversations"].insert_one(
|
||||
{
|
||||
"_id": "cccccccccccccccccccccccc",
|
||||
"user": "alice",
|
||||
"queries": [
|
||||
{
|
||||
"prompt": "hi",
|
||||
# Unknown attachment objectid — not present in PG.
|
||||
"attachments": ["bbbbbbbbbbbbbbbbbbbbbbb0"],
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
stats = _backfill_conversations(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
assert stats["unresolved_attachment_refs"] == 1
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT attachments FROM conversation_messages "
|
||||
"WHERE position = 0"
|
||||
)
|
||||
).one()._mapping
|
||||
assert list(row["attachments"]) == []
|
||||
|
||||
def test_conversations_rerun_does_not_double_messages(
|
||||
self, pg_conn, mongo_db
|
||||
):
|
||||
mongo_db["conversations"].insert_one(
|
||||
{
|
||||
"_id": "cccccccccccccccccccccccc",
|
||||
"user": "alice",
|
||||
"queries": [
|
||||
{"prompt": "a", "response": "1"},
|
||||
{"prompt": "b", "response": "2"},
|
||||
],
|
||||
}
|
||||
)
|
||||
_backfill_conversations(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
_backfill_conversations(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
conv_count = pg_conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM conversations "
|
||||
"WHERE legacy_mongo_id = 'cccccccccccccccccccccccc'"
|
||||
)
|
||||
).scalar()
|
||||
msg_count = pg_conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM conversation_messages "
|
||||
"WHERE conversation_id = ("
|
||||
" SELECT id FROM conversations "
|
||||
" WHERE legacy_mongo_id = 'cccccccccccccccccccccccc'"
|
||||
")"
|
||||
)
|
||||
).scalar()
|
||||
assert conv_count == 1
|
||||
assert msg_count == 2
|
||||
|
||||
def test_conversations_rerun_truncates_removed_tail_messages(
|
||||
self, pg_conn, mongo_db
|
||||
):
|
||||
# First run: 3 messages. Then Mongo truncates to 1. Second run
|
||||
# should drop positions 1 and 2.
|
||||
mongo_db["conversations"].insert_one(
|
||||
{
|
||||
"_id": "cccccccccccccccccccccccc",
|
||||
"user": "alice",
|
||||
"queries": [
|
||||
{"prompt": "a", "response": "1"},
|
||||
{"prompt": "b", "response": "2"},
|
||||
{"prompt": "c", "response": "3"},
|
||||
],
|
||||
}
|
||||
)
|
||||
_backfill_conversations(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
mongo_db["conversations"].update_one(
|
||||
{"_id": "cccccccccccccccccccccccc"},
|
||||
{"$set": {"queries": [{"prompt": "a", "response": "1-updated"}]}},
|
||||
)
|
||||
_backfill_conversations(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
rows = pg_conn.execute(
|
||||
text(
|
||||
"SELECT position, response FROM conversation_messages "
|
||||
"WHERE conversation_id = ("
|
||||
" SELECT id FROM conversations "
|
||||
" WHERE legacy_mongo_id = 'cccccccccccccccccccccccc'"
|
||||
") ORDER BY position"
|
||||
)
|
||||
).fetchall()
|
||||
assert len(rows) == 1
|
||||
assert rows[0]._mapping["response"] == "1-updated"
|
||||
@@ -1,144 +0,0 @@
|
||||
"""Integration tests for ``_backfill_sources``.
|
||||
|
||||
The notable translation here is the optional Mongo ``user`` field: system
|
||||
seed rows arrive with no ``user`` (or ``user="system"``) and must land
|
||||
with ``user_id = '__system__'`` in Postgres so the NOT NULL constraint
|
||||
accepts them.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import mongomock
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[3]))
|
||||
|
||||
from scripts.db.backfill import ( # noqa: E402
|
||||
SYSTEM_USER_ID,
|
||||
_backfill_sources,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mongo_db() -> Any:
|
||||
client = mongomock.MongoClient()
|
||||
return client["docsgpt_test"]
|
||||
|
||||
|
||||
class TestBackfillSources:
|
||||
def test_sources_system_rows_get_sentinel_user(self, pg_conn, mongo_db):
|
||||
mongo_db["sources"].insert_many(
|
||||
[
|
||||
{
|
||||
"_id": "555555555555555555555555",
|
||||
"name": "Seed Source",
|
||||
"type": "url",
|
||||
# Intentionally no ``user`` field.
|
||||
},
|
||||
{
|
||||
"_id": "666666666666666666666666",
|
||||
"user": "alice",
|
||||
"name": "Alice's Upload",
|
||||
"type": "file",
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
stats = _backfill_sources(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
assert stats["seen"] == 2
|
||||
assert stats["written"] == 2
|
||||
|
||||
rows = pg_conn.execute(
|
||||
text(
|
||||
"SELECT user_id, name FROM sources "
|
||||
"WHERE legacy_mongo_id IN "
|
||||
"('555555555555555555555555', '666666666666666666666666') "
|
||||
"ORDER BY name"
|
||||
)
|
||||
).fetchall()
|
||||
assert len(rows) == 2
|
||||
# Alphabetical by name: "Alice's Upload" comes before "Seed Source".
|
||||
assert rows[0]._mapping["user_id"] == "alice"
|
||||
assert rows[1]._mapping["user_id"] == SYSTEM_USER_ID
|
||||
|
||||
def test_sources_preserves_legacy_fields_under_metadata(
|
||||
self, pg_conn, mongo_db
|
||||
):
|
||||
mongo_db["sources"].insert_one(
|
||||
{
|
||||
"_id": "555555555555555555555555",
|
||||
"user": "alice",
|
||||
"name": "s1",
|
||||
"type": "url",
|
||||
"status": "ingested", # legacy ingestion field
|
||||
"reason": None,
|
||||
}
|
||||
)
|
||||
_backfill_sources(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
meta = pg_conn.execute(
|
||||
text(
|
||||
"SELECT metadata FROM sources "
|
||||
"WHERE legacy_mongo_id = '555555555555555555555555'"
|
||||
)
|
||||
).scalar()
|
||||
assert meta["legacy_fields"]["status"] == "ingested"
|
||||
|
||||
def test_sources_rerun_does_not_duplicate(self, pg_conn, mongo_db):
|
||||
mongo_db["sources"].insert_one(
|
||||
{
|
||||
"_id": "555555555555555555555555",
|
||||
"user": "alice",
|
||||
"name": "s1",
|
||||
"type": "url",
|
||||
}
|
||||
)
|
||||
_backfill_sources(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
_backfill_sources(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
count = pg_conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM sources "
|
||||
"WHERE legacy_mongo_id = '555555555555555555555555'"
|
||||
)
|
||||
).scalar()
|
||||
assert count == 1
|
||||
|
||||
def test_sources_rerun_converges_on_name(self, pg_conn, mongo_db):
|
||||
mongo_db["sources"].insert_one(
|
||||
{
|
||||
"_id": "555555555555555555555555",
|
||||
"user": "alice",
|
||||
"name": "old-name",
|
||||
"type": "url",
|
||||
}
|
||||
)
|
||||
_backfill_sources(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
mongo_db["sources"].update_one(
|
||||
{"_id": "555555555555555555555555"},
|
||||
{"$set": {"name": "new-name"}},
|
||||
)
|
||||
_backfill_sources(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
name = pg_conn.execute(
|
||||
text(
|
||||
"SELECT name FROM sources "
|
||||
"WHERE legacy_mongo_id = '555555555555555555555555'"
|
||||
)
|
||||
).scalar()
|
||||
assert name == "new-name"
|
||||
@@ -1,380 +0,0 @@
|
||||
"""Integration tests for the per-tool-child collections: todos, notes,
|
||||
and memories.
|
||||
|
||||
These all depend on ``user_tools`` having been backfilled first, because
|
||||
:func:`scripts.db.backfill._build_tool_id_map` joins Mongo
|
||||
``user_tools._id`` to Postgres ``user_tools.id`` via ``(user_id, name)``.
|
||||
Each test seeds a single tool row to keep fixtures minimal.
|
||||
|
||||
Memories are the odd one out: the SQL uses ``ON CONFLICT DO NOTHING`` and
|
||||
the table has no ``legacy_mongo_id`` column, so a re-run can't converge
|
||||
on mutated content. The negative test locks in the current behavior so a
|
||||
future fix (adding ``legacy_mongo_id`` + DO UPDATE, tracked in
|
||||
``migration-postgres.md``) will flip the assertion.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import mongomock
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[3]))
|
||||
|
||||
from scripts.db.backfill import ( # noqa: E402
|
||||
_backfill_memories,
|
||||
_backfill_notes,
|
||||
_backfill_todos,
|
||||
_backfill_user_tools,
|
||||
)
|
||||
|
||||
|
||||
_TOOL_MONGO_ID = "507f1f77bcf86cd799439011"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mongo_db() -> Any:
|
||||
client = mongomock.MongoClient()
|
||||
return client["docsgpt_test"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def seeded_tool(pg_conn, mongo_db) -> str:
|
||||
"""Seed one tool in both Mongo and Postgres and return its PG UUID.
|
||||
|
||||
Having a user_tools row on both sides is a prerequisite for every
|
||||
backfill function in this file — ``_build_tool_id_map`` needs it to
|
||||
resolve ``tool_id`` FKs.
|
||||
"""
|
||||
mongo_db["user_tools"].insert_one(
|
||||
{
|
||||
"_id": _TOOL_MONGO_ID,
|
||||
"user": "alice",
|
||||
"name": "my-tool",
|
||||
}
|
||||
)
|
||||
_backfill_user_tools(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
return str(
|
||||
pg_conn.execute(
|
||||
text("SELECT id FROM user_tools WHERE legacy_mongo_id = :lid"),
|
||||
{"lid": _TOOL_MONGO_ID},
|
||||
).scalar()
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# todos
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBackfillTodos:
|
||||
def test_todos_status_completed_translation(self, pg_conn, mongo_db, seeded_tool):
|
||||
# Two todos: one "open", one "completed". Asserts the Mongo
|
||||
# ``status`` string is translated to the PG ``completed`` bool.
|
||||
mongo_db["todos"].insert_many(
|
||||
[
|
||||
{
|
||||
"_id": "111111111111111111111111",
|
||||
"user_id": "alice",
|
||||
"tool_id": _TOOL_MONGO_ID,
|
||||
"todo_id": 1,
|
||||
"title": "Write tests",
|
||||
"status": "open",
|
||||
},
|
||||
{
|
||||
"_id": "222222222222222222222222",
|
||||
"user_id": "alice",
|
||||
"tool_id": _TOOL_MONGO_ID,
|
||||
"todo_id": 2,
|
||||
"title": "Ship it",
|
||||
"status": "completed",
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
stats = _backfill_todos(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
assert stats["seen"] == 2
|
||||
assert stats["written"] == 2
|
||||
|
||||
rows = pg_conn.execute(
|
||||
text(
|
||||
"SELECT todo_id, title, completed, legacy_mongo_id FROM todos "
|
||||
"WHERE user_id = 'alice' ORDER BY todo_id"
|
||||
)
|
||||
).fetchall()
|
||||
assert len(rows) == 2
|
||||
assert rows[0]._mapping["todo_id"] == 1
|
||||
assert rows[0]._mapping["completed"] is False
|
||||
assert rows[1]._mapping["todo_id"] == 2
|
||||
assert rows[1]._mapping["completed"] is True
|
||||
# legacy_mongo_id preserved for idempotency on re-run.
|
||||
assert rows[0]._mapping["legacy_mongo_id"] == "111111111111111111111111"
|
||||
|
||||
def test_todos_legacy_fields_stashed_in_metadata(
|
||||
self, pg_conn, mongo_db, seeded_tool
|
||||
):
|
||||
# Legacy top-level ``conversation_id`` field should survive under
|
||||
# metadata.legacy_fields.
|
||||
mongo_db["todos"].insert_one(
|
||||
{
|
||||
"_id": "111111111111111111111111",
|
||||
"user_id": "alice",
|
||||
"tool_id": _TOOL_MONGO_ID,
|
||||
"title": "t",
|
||||
"status": "open",
|
||||
"conversation_id": "abc-legacy",
|
||||
}
|
||||
)
|
||||
_backfill_todos(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
meta = pg_conn.execute(
|
||||
text("SELECT metadata FROM todos WHERE legacy_mongo_id = :lid"),
|
||||
{"lid": "111111111111111111111111"},
|
||||
).scalar()
|
||||
assert meta["legacy_fields"]["conversation_id"] == "abc-legacy"
|
||||
|
||||
def test_todos_rerun_does_not_duplicate(self, pg_conn, mongo_db, seeded_tool):
|
||||
mongo_db["todos"].insert_one(
|
||||
{
|
||||
"_id": "111111111111111111111111",
|
||||
"user_id": "alice",
|
||||
"tool_id": _TOOL_MONGO_ID,
|
||||
"todo_id": 1,
|
||||
"title": "t",
|
||||
"status": "open",
|
||||
}
|
||||
)
|
||||
_backfill_todos(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
_backfill_todos(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
count = pg_conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM todos "
|
||||
"WHERE legacy_mongo_id = '111111111111111111111111'"
|
||||
)
|
||||
).scalar()
|
||||
assert count == 1
|
||||
|
||||
def test_todos_rerun_converges_on_status_and_title(
|
||||
self, pg_conn, mongo_db, seeded_tool
|
||||
):
|
||||
mongo_db["todos"].insert_one(
|
||||
{
|
||||
"_id": "111111111111111111111111",
|
||||
"user_id": "alice",
|
||||
"tool_id": _TOOL_MONGO_ID,
|
||||
"todo_id": 1,
|
||||
"title": "Old title",
|
||||
"status": "open",
|
||||
}
|
||||
)
|
||||
_backfill_todos(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
mongo_db["todos"].update_one(
|
||||
{"_id": "111111111111111111111111"},
|
||||
{"$set": {"status": "completed", "title": "Done title"}},
|
||||
)
|
||||
_backfill_todos(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT title, completed FROM todos "
|
||||
"WHERE legacy_mongo_id = '111111111111111111111111'"
|
||||
)
|
||||
).one()._mapping
|
||||
assert row["title"] == "Done title"
|
||||
assert row["completed"] is True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# notes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBackfillNotes:
|
||||
def test_notes_translates_note_to_content(
|
||||
self, pg_conn, mongo_db, seeded_tool
|
||||
):
|
||||
mongo_db["notes"].insert_one(
|
||||
{
|
||||
"_id": "333333333333333333333333",
|
||||
"user_id": "alice",
|
||||
"tool_id": _TOOL_MONGO_ID,
|
||||
"note": "hello world",
|
||||
# No explicit title — backfill should fall back to "note".
|
||||
}
|
||||
)
|
||||
_backfill_notes(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT title, content, legacy_mongo_id FROM notes "
|
||||
"WHERE user_id = 'alice'"
|
||||
)
|
||||
).one()._mapping
|
||||
assert row["content"] == "hello world"
|
||||
assert row["title"] == "note"
|
||||
# legacy_mongo_id column landed recently — lock in its presence.
|
||||
assert row["legacy_mongo_id"] == "333333333333333333333333"
|
||||
|
||||
def test_notes_uses_title_when_present(self, pg_conn, mongo_db, seeded_tool):
|
||||
mongo_db["notes"].insert_one(
|
||||
{
|
||||
"_id": "333333333333333333333333",
|
||||
"user_id": "alice",
|
||||
"tool_id": _TOOL_MONGO_ID,
|
||||
"title": "My title",
|
||||
"content": "direct content",
|
||||
}
|
||||
)
|
||||
_backfill_notes(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
row = pg_conn.execute(
|
||||
text("SELECT title, content FROM notes WHERE user_id = 'alice'")
|
||||
).one()._mapping
|
||||
assert row["title"] == "My title"
|
||||
assert row["content"] == "direct content"
|
||||
|
||||
def test_notes_rerun_converges_on_content(
|
||||
self, pg_conn, mongo_db, seeded_tool
|
||||
):
|
||||
mongo_db["notes"].insert_one(
|
||||
{
|
||||
"_id": "333333333333333333333333",
|
||||
"user_id": "alice",
|
||||
"tool_id": _TOOL_MONGO_ID,
|
||||
"note": "v1",
|
||||
}
|
||||
)
|
||||
_backfill_notes(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
mongo_db["notes"].update_one(
|
||||
{"_id": "333333333333333333333333"},
|
||||
{"$set": {"note": "v2"}},
|
||||
)
|
||||
_backfill_notes(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
rows = pg_conn.execute(
|
||||
text(
|
||||
"SELECT content FROM notes "
|
||||
"WHERE legacy_mongo_id = '333333333333333333333333'"
|
||||
)
|
||||
).fetchall()
|
||||
assert len(rows) == 1
|
||||
assert rows[0]._mapping["content"] == "v2"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# memories
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBackfillMemories:
|
||||
def test_memories_happy_shape(self, pg_conn, mongo_db, seeded_tool):
|
||||
mongo_db["memories"].insert_one(
|
||||
{
|
||||
"_id": "444444444444444444444444",
|
||||
"user_id": "alice",
|
||||
"tool_id": _TOOL_MONGO_ID,
|
||||
"path": "/foo",
|
||||
"content": "memory body",
|
||||
}
|
||||
)
|
||||
_backfill_memories(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
row = pg_conn.execute(
|
||||
text("SELECT path, content FROM memories WHERE user_id = 'alice'")
|
||||
).one()._mapping
|
||||
assert row["path"] == "/foo"
|
||||
assert row["content"] == "memory body"
|
||||
|
||||
def test_memories_rerun_does_not_duplicate(
|
||||
self, pg_conn, mongo_db, seeded_tool
|
||||
):
|
||||
# The unique index is (user_id, tool_id, path). Same-row re-insert
|
||||
# is expected to be a no-op because the SQL uses DO NOTHING.
|
||||
mongo_db["memories"].insert_one(
|
||||
{
|
||||
"_id": "444444444444444444444444",
|
||||
"user_id": "alice",
|
||||
"tool_id": _TOOL_MONGO_ID,
|
||||
"path": "/foo",
|
||||
"content": "v1",
|
||||
}
|
||||
)
|
||||
_backfill_memories(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
_backfill_memories(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
count = pg_conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM memories "
|
||||
"WHERE user_id = 'alice' AND path = '/foo'"
|
||||
)
|
||||
).scalar()
|
||||
assert count == 1
|
||||
|
||||
def test_memories_rerun_does_not_converge_content(
|
||||
self, pg_conn, mongo_db, seeded_tool
|
||||
):
|
||||
# KNOWN non-idempotent behavior: memories have no legacy_mongo_id
|
||||
# column and the SQL uses ``ON CONFLICT DO NOTHING`` rather than
|
||||
# DO UPDATE. A content change in Mongo is NOT reflected in PG on
|
||||
# re-run. See migration-postgres.md for the tracked fix. When the
|
||||
# backfill gains a legacy_mongo_id column + DO UPDATE branch,
|
||||
# flip this assertion to assert the new content wins.
|
||||
mongo_db["memories"].insert_one(
|
||||
{
|
||||
"_id": "444444444444444444444444",
|
||||
"user_id": "alice",
|
||||
"tool_id": _TOOL_MONGO_ID,
|
||||
"path": "/foo",
|
||||
"content": "original",
|
||||
}
|
||||
)
|
||||
_backfill_memories(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
mongo_db["memories"].update_one(
|
||||
{"_id": "444444444444444444444444"},
|
||||
{"$set": {"content": "updated"}},
|
||||
)
|
||||
_backfill_memories(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
content = pg_conn.execute(
|
||||
text(
|
||||
"SELECT content FROM memories "
|
||||
"WHERE user_id = 'alice' AND path = '/foo'"
|
||||
)
|
||||
).scalar()
|
||||
# Lock in the non-convergent behavior.
|
||||
assert content == "original"
|
||||
@@ -1,348 +0,0 @@
|
||||
"""Integration tests for ``_backfill_users`` / ``_backfill_prompts`` /
|
||||
``_backfill_user_tools`` against an ephemeral Postgres.
|
||||
|
||||
These exercise the Mongo→PG shape translations end-to-end: a fake Mongo
|
||||
(mongomock) is populated with representative docs, the ``_backfill_*``
|
||||
function under test runs against ``pg_conn``, and the resulting rows are
|
||||
asserted to carry the translated fields. Each function is also run twice
|
||||
to validate idempotency (no row duplication) and — where the SQL uses
|
||||
``DO UPDATE`` — convergence on mutated source data.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import mongomock
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
# Make ``scripts.db.backfill`` importable (scripts/ isn't on sys.path by default).
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[3]))
|
||||
|
||||
from scripts.db.backfill import ( # noqa: E402
|
||||
SYSTEM_USER_ID,
|
||||
_backfill_prompts,
|
||||
_backfill_user_tools,
|
||||
_backfill_users,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mongo_db() -> Any:
|
||||
"""Fresh in-memory Mongo database per test."""
|
||||
client = mongomock.MongoClient()
|
||||
return client["docsgpt_test"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# users
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBackfillUsers:
|
||||
def test_users_happy_shape_merges_agent_preferences(self, pg_conn, mongo_db):
|
||||
mongo_db["users"].insert_many(
|
||||
[
|
||||
{
|
||||
"user_id": "alice",
|
||||
"agent_preferences": {
|
||||
"pinned": ["agent-1"],
|
||||
"theme": "dark",
|
||||
},
|
||||
},
|
||||
{
|
||||
"user_id": "bob",
|
||||
"agent_preferences": {},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
stats = _backfill_users(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
assert stats["seen"] == 2
|
||||
assert stats["written"] == 2
|
||||
assert stats["skipped_no_user_id"] == 0
|
||||
|
||||
rows = pg_conn.execute(
|
||||
text(
|
||||
"SELECT user_id, agent_preferences FROM users "
|
||||
"WHERE user_id IN ('alice', 'bob') ORDER BY user_id"
|
||||
)
|
||||
).fetchall()
|
||||
assert len(rows) == 2
|
||||
|
||||
alice = rows[0]._mapping
|
||||
bob = rows[1]._mapping
|
||||
# Unknown top-level prefs (``theme``) survive untouched.
|
||||
assert alice["agent_preferences"]["theme"] == "dark"
|
||||
assert alice["agent_preferences"]["pinned"] == ["agent-1"]
|
||||
# Missing ``shared_with_me`` gets filled to [].
|
||||
assert alice["agent_preferences"]["shared_with_me"] == []
|
||||
assert bob["agent_preferences"]["pinned"] == []
|
||||
assert bob["agent_preferences"]["shared_with_me"] == []
|
||||
|
||||
def test_users_skips_rows_without_user_id(self, pg_conn, mongo_db):
|
||||
mongo_db["users"].insert_many(
|
||||
[{"user_id": "alice"}, {"agent_preferences": {"pinned": []}}]
|
||||
)
|
||||
|
||||
stats = _backfill_users(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
assert stats["seen"] == 2
|
||||
assert stats["skipped_no_user_id"] == 1
|
||||
assert stats["written"] == 1
|
||||
|
||||
def test_users_rerun_does_not_duplicate(self, pg_conn, mongo_db):
|
||||
mongo_db["users"].insert_one(
|
||||
{"user_id": "alice", "agent_preferences": {"pinned": ["a-1"]}}
|
||||
)
|
||||
|
||||
_backfill_users(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
_backfill_users(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
count = pg_conn.execute(
|
||||
text("SELECT count(*) FROM users WHERE user_id = 'alice'")
|
||||
).scalar()
|
||||
assert count == 1
|
||||
|
||||
def test_users_rerun_converges_on_mutated_pinned(self, pg_conn, mongo_db):
|
||||
mongo_db["users"].insert_one(
|
||||
{"user_id": "alice", "agent_preferences": {"pinned": ["a-1"]}}
|
||||
)
|
||||
_backfill_users(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
# Mutate Mongo.
|
||||
mongo_db["users"].update_one(
|
||||
{"user_id": "alice"},
|
||||
{"$set": {"agent_preferences": {"pinned": ["a-1", "a-2"]}}},
|
||||
)
|
||||
_backfill_users(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
prefs = pg_conn.execute(
|
||||
text("SELECT agent_preferences FROM users WHERE user_id = 'alice'")
|
||||
).scalar()
|
||||
assert prefs["pinned"] == ["a-1", "a-2"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# prompts
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBackfillPrompts:
|
||||
def test_prompts_happy_shape(self, pg_conn, mongo_db):
|
||||
mongo_db["prompts"].insert_many(
|
||||
[
|
||||
{
|
||||
"_id": "507f1f77bcf86cd799439011",
|
||||
"user": "alice",
|
||||
"name": "Greet",
|
||||
"content": "Say hi",
|
||||
},
|
||||
{
|
||||
"_id": "507f1f77bcf86cd799439012",
|
||||
"user": "system",
|
||||
"name": "Template",
|
||||
"content": "Seed",
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
stats = _backfill_prompts(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
assert stats["seen"] == 2
|
||||
assert stats["written"] == 2
|
||||
|
||||
rows = pg_conn.execute(
|
||||
text(
|
||||
"SELECT user_id, name, content, legacy_mongo_id "
|
||||
"FROM prompts ORDER BY name"
|
||||
)
|
||||
).fetchall()
|
||||
assert len(rows) == 2
|
||||
|
||||
greet = rows[0]._mapping
|
||||
template = rows[1]._mapping
|
||||
assert greet["user_id"] == "alice"
|
||||
assert greet["name"] == "Greet"
|
||||
assert greet["content"] == "Say hi"
|
||||
assert greet["legacy_mongo_id"] == "507f1f77bcf86cd799439011"
|
||||
# Legacy ``user="system"`` collapses to the sentinel.
|
||||
assert template["user_id"] == SYSTEM_USER_ID
|
||||
|
||||
def test_prompts_rerun_does_not_duplicate(self, pg_conn, mongo_db):
|
||||
mongo_db["prompts"].insert_one(
|
||||
{
|
||||
"_id": "507f1f77bcf86cd799439011",
|
||||
"user": "alice",
|
||||
"name": "Greet",
|
||||
"content": "Say hi",
|
||||
}
|
||||
)
|
||||
_backfill_prompts(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
_backfill_prompts(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
count = pg_conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM prompts WHERE "
|
||||
"legacy_mongo_id = '507f1f77bcf86cd799439011'"
|
||||
)
|
||||
).scalar()
|
||||
assert count == 1
|
||||
|
||||
def test_prompts_rerun_converges_on_content(self, pg_conn, mongo_db):
|
||||
mongo_db["prompts"].insert_one(
|
||||
{
|
||||
"_id": "507f1f77bcf86cd799439011",
|
||||
"user": "alice",
|
||||
"name": "Greet",
|
||||
"content": "v1",
|
||||
}
|
||||
)
|
||||
_backfill_prompts(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
mongo_db["prompts"].update_one(
|
||||
{"_id": "507f1f77bcf86cd799439011"},
|
||||
{"$set": {"content": "v2"}},
|
||||
)
|
||||
_backfill_prompts(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT content FROM prompts WHERE "
|
||||
"legacy_mongo_id = '507f1f77bcf86cd799439011'"
|
||||
)
|
||||
).scalar()
|
||||
assert row == "v2"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# user_tools
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBackfillUserTools:
|
||||
def test_user_tools_nested_values_stringified(self, pg_conn, mongo_db):
|
||||
# Use a dict with a datetime-like / non-JSON-native type nested
|
||||
# inside config to verify ``default=str`` gets applied.
|
||||
mongo_db["user_tools"].insert_one(
|
||||
{
|
||||
"_id": "507f1f77bcf86cd799439011",
|
||||
"user": "alice",
|
||||
"name": "calendar",
|
||||
"displayName": "Calendar",
|
||||
"customName": "My Calendar",
|
||||
"description": "cal desc",
|
||||
"config": {
|
||||
"api_key": "secret",
|
||||
"expires_at": datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||||
},
|
||||
"configRequirements": {"api_key": {"type": "string"}},
|
||||
"actions": [{"name": "create_event"}],
|
||||
"status": True,
|
||||
}
|
||||
)
|
||||
|
||||
_backfill_user_tools(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT name, custom_name, display_name, description, "
|
||||
"config, config_requirements, actions, status "
|
||||
"FROM user_tools WHERE legacy_mongo_id = :lid"
|
||||
),
|
||||
{"lid": "507f1f77bcf86cd799439011"},
|
||||
).one()._mapping
|
||||
|
||||
assert row["name"] == "calendar"
|
||||
assert row["custom_name"] == "My Calendar"
|
||||
assert row["display_name"] == "Calendar"
|
||||
assert row["description"] == "cal desc"
|
||||
assert row["status"] is True
|
||||
# Nested config datetime is serialized to a string (default=str).
|
||||
assert row["config"]["api_key"] == "secret"
|
||||
assert isinstance(row["config"]["expires_at"], str)
|
||||
assert row["config_requirements"] == {"api_key": {"type": "string"}}
|
||||
assert row["actions"] == [{"name": "create_event"}]
|
||||
|
||||
def test_user_tools_rerun_converges(self, pg_conn, mongo_db):
|
||||
mongo_db["user_tools"].insert_one(
|
||||
{
|
||||
"_id": "507f1f77bcf86cd799439011",
|
||||
"user": "alice",
|
||||
"name": "cal",
|
||||
"config": {"k": "v1"},
|
||||
}
|
||||
)
|
||||
_backfill_user_tools(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
mongo_db["user_tools"].update_one(
|
||||
{"_id": "507f1f77bcf86cd799439011"},
|
||||
{"$set": {"config": {"k": "v2"}, "description": "new"}},
|
||||
)
|
||||
_backfill_user_tools(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
|
||||
rows = pg_conn.execute(
|
||||
text(
|
||||
"SELECT config, description FROM user_tools "
|
||||
"WHERE legacy_mongo_id = '507f1f77bcf86cd799439011'"
|
||||
)
|
||||
).fetchall()
|
||||
assert len(rows) == 1
|
||||
assert rows[0]._mapping["config"] == {"k": "v2"}
|
||||
assert rows[0]._mapping["description"] == "new"
|
||||
|
||||
def test_user_tools_rerun_does_not_duplicate(self, pg_conn, mongo_db):
|
||||
mongo_db["user_tools"].insert_one(
|
||||
{
|
||||
"_id": "507f1f77bcf86cd799439011",
|
||||
"user": "alice",
|
||||
"name": "cal",
|
||||
}
|
||||
)
|
||||
_backfill_user_tools(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
_backfill_user_tools(
|
||||
conn=pg_conn, mongo_db=mongo_db, batch_size=100, dry_run=False
|
||||
)
|
||||
count = pg_conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM user_tools "
|
||||
"WHERE legacy_mongo_id = '507f1f77bcf86cd799439011'"
|
||||
)
|
||||
).scalar()
|
||||
assert count == 1
|
||||
|
||||
|
||||
Reference in new issue
Block a user