diff --git a/tests/storage/db/test_backfill_integration_connector_sessions.py b/tests/storage/db/test_backfill_integration_connector_sessions.py deleted file mode 100644 index 5222e488..00000000 --- a/tests/storage/db/test_backfill_integration_connector_sessions.py +++ /dev/null @@ -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 diff --git a/tests/storage/db/test_backfill_integration_conversations.py b/tests/storage/db/test_backfill_integration_conversations.py deleted file mode 100644 index 6e9d792e..00000000 --- a/tests/storage/db/test_backfill_integration_conversations.py +++ /dev/null @@ -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" diff --git a/tests/storage/db/test_backfill_integration_sources.py b/tests/storage/db/test_backfill_integration_sources.py deleted file mode 100644 index c09bd70e..00000000 --- a/tests/storage/db/test_backfill_integration_sources.py +++ /dev/null @@ -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" diff --git a/tests/storage/db/test_backfill_integration_todos_notes_memories.py b/tests/storage/db/test_backfill_integration_todos_notes_memories.py deleted file mode 100644 index 7edd04f5..00000000 --- a/tests/storage/db/test_backfill_integration_todos_notes_memories.py +++ /dev/null @@ -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" diff --git a/tests/storage/db/test_backfill_integration_users.py b/tests/storage/db/test_backfill_integration_users.py deleted file mode 100644 index 0958cf1a..00000000 --- a/tests/storage/db/test_backfill_integration_users.py +++ /dev/null @@ -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 - -