diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5ba50470..7cc0de63 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -65,12 +65,31 @@ repos: # Run main application tests - id: pytest-main name: pytest (main app) - entry: uv run pytest tests/ + entry: uv run pytest tests/ --ignore=tests/alembic/ language: system files: ^(src/|tests/).*\.py$ stages: [pre-push] pass_filenames: false + # Run Alembic tests only when migrations change + - id: pytest-alembic + name: pytest (alembic migrations) + entry: uv run pytest tests/alembic/ + language: system + files: ^(migrations/.*\.py|tests/alembic/.*\.py)$ + stages: [pre-push] + pass_filenames: false + require_serial: true + + # Ensure each alembic migration revision has a corresponding test file + - id: ensure-alembic-coverage + name: ensure alembic migration test coverage + entry: uv run python scripts/ensure_alembic_tests.py + language: system + files: ^(migrations/versions/.*\.py|tests/alembic/revisions/.*\.py)$ + stages: [pre-push] + pass_filenames: false + # Run Python SDK tests (if they exist) - id: pytest-python-sdk name: pytest (Python SDK) diff --git a/migrations/versions/08894082221a_replace_collection_name_with_observer_.py b/migrations/versions/08894082221a_replace_collection_name_with_observer_.py index cadad00f..87067d67 100644 --- a/migrations/versions/08894082221a_replace_collection_name_with_observer_.py +++ b/migrations/versions/08894082221a_replace_collection_name_with_observer_.py @@ -476,12 +476,10 @@ def downgrade() -> None: # Make peer_name NOT NULL op.alter_column("collections", "peer_name", nullable=False, schema=schema) - # Recreate the foreign key constraint for peer_name - if not fk_exists( - "collections", "collections_peer_name_workspace_name_fkey", inspector - ): + # Recreate the legacy foreign key constraint for peer_name + if not fk_exists("collections", "fk_collections_peer_name_peers", inspector): op.create_foreign_key( - "collections_peer_name_workspace_name_fkey", + "fk_collections_peer_name_peers", "collections", "peers", ["peer_name", "workspace_name"], @@ -530,7 +528,7 @@ def downgrade() -> None: # Step 6: Make documents collection_name NOT NULL op.alter_column("documents", "collection_name", nullable=False, schema=schema) - # Step 6a: Restore peer_name column to documents (set to observed value) + # Step 6a: Restore peer_name column to documents (set to observer value) if not column_exists("documents", "peer_name", inspector): op.add_column( "documents", @@ -637,14 +635,14 @@ def downgrade() -> None: schema=schema, ) - # Step 11: Recreate the old foreign key constraint from documents to collections + # Step 11: Recreate the legacy foreign key constraint from documents to collections if not fk_exists( "documents", - "documents_collection_name_peer_name_workspace_name_fkey", + "fk_documents_collection_name_collections", inspector, ): op.create_foreign_key( - "documents_collection_name_peer_name_workspace_name_fkey", + "fk_documents_collection_name_collections", "documents", "collections", ["collection_name", "peer_name", "workspace_name"], diff --git a/migrations/versions/556a16564f50_add_user_id_and_app_id_to_tables.py b/migrations/versions/556a16564f50_add_user_id_and_app_id_to_tables.py index d766f124..05338d0e 100644 --- a/migrations/versions/556a16564f50_add_user_id_and_app_id_to_tables.py +++ b/migrations/versions/556a16564f50_add_user_id_and_app_id_to_tables.py @@ -435,18 +435,18 @@ def downgrade(): op.alter_column("documents", "user_id", nullable=True, schema=schema) print("Made app_id and user_id nullable again for documents table") - # Drop foreign key constraints - try: + # Drop foreign key constraints (guarded) + if fk_exists("documents", "documents_user_id_fkey", inspector): op.drop_constraint("documents_user_id_fkey", "documents", schema=schema) print("Dropped user_id foreign key for documents table") - except Exception as e: - print(f"Error dropping user_id foreign key for documents table: {e}") + else: + print("user_id foreign key for documents table does not exist; skipping drop") - try: + if fk_exists("documents", "documents_app_id_fkey", inspector): op.drop_constraint("documents_app_id_fkey", "documents", schema=schema) print("Dropped app_id foreign key for documents table") - except Exception as e: - print(f"Error dropping app_id foreign key for documents table: {e}") + else: + print("app_id foreign key for documents table does not exist; skipping drop") # Drop the columns op.drop_column("documents", "user_id", schema=schema) @@ -464,12 +464,12 @@ def downgrade(): op.alter_column("collections", "app_id", nullable=True, schema=schema) print("Made app_id nullable again for collections table") - # Drop foreign key constraint - try: + # Drop foreign key constraint (guarded) + if fk_exists("collections", "collections_app_id_fkey", inspector): op.drop_constraint("collections_app_id_fkey", "collections", schema=schema) print("Dropped app_id foreign key for collections table") - except Exception as e: - print(f"Error dropping app_id foreign key for collections table: {e}") + else: + print("app_id foreign key for collections table does not exist; skipping drop") # Drop the column op.drop_column("collections", "app_id", schema=schema) @@ -484,20 +484,26 @@ def downgrade(): ) print("Dropped index ix_metamessages_app_id") - # Make app_id nullable again - op.alter_column("metamessages", "app_id", nullable=True, schema=schema) - print("Made app_id nullable again for metamessages table") + # Make app_id nullable again (only if column exists) + if column_exists("metamessages", "app_id", inspector): + op.alter_column("metamessages", "app_id", nullable=True, schema=schema) + print("Made app_id nullable again for metamessages table") + else: + print("metamessages.app_id column does not exist; skipping alter nullable") - # Drop foreign key constraint - try: + # Drop foreign key constraint (guarded) + if fk_exists("metamessages", "metamessages_app_id_fkey", inspector): op.drop_constraint("metamessages_app_id_fkey", "metamessages", schema=schema) print("Dropped app_id foreign key for metamessages table") - except Exception as e: - print(f"Error dropping app_id foreign key for metamessages table: {e}") + else: + print("app_id foreign key for metamessages table does not exist; skipping drop") - # Drop the column - op.drop_column("metamessages", "app_id", schema=schema) - print("Dropped app_id column from metamessages table") + # Drop the column (only if present) + if column_exists("metamessages", "app_id", inspector): + op.drop_column("metamessages", "app_id", schema=schema) + print("Dropped app_id column from metamessages table") + else: + print("metamessages.app_id column does not exist; skipping drop column") # 2. Messages table @@ -514,18 +520,18 @@ def downgrade(): op.alter_column("messages", "user_id", nullable=True, schema=schema) print("Made app_id and user_id nullable again for messages table") - # Drop foreign key constraints - try: + # Drop foreign key constraints (guarded) + if fk_exists("messages", "messages_user_id_fkey", inspector): op.drop_constraint("messages_user_id_fkey", "messages", schema=schema) print("Dropped user_id foreign key for messages table") - except Exception as e: - print(f"Error dropping user_id foreign key for messages table: {e}") + else: + print("user_id foreign key for messages table does not exist; skipping drop") - try: + if fk_exists("messages", "messages_app_id_fkey", inspector): op.drop_constraint("messages_app_id_fkey", "messages", schema=schema) print("Dropped app_id foreign key for messages table") - except Exception as e: - print(f"Error dropping app_id foreign key for messages table: {e}") + else: + print("app_id foreign key for messages table does not exist; skipping drop") # Drop the columns op.drop_column("messages", "user_id", schema=schema) @@ -543,12 +549,12 @@ def downgrade(): op.alter_column("sessions", "app_id", nullable=True, schema=schema) print("Made app_id nullable again for sessions table") - # Drop foreign key constraint - try: + # Drop foreign key constraint (guarded) + if fk_exists("sessions", "sessions_app_id_fkey", inspector): op.drop_constraint("sessions_app_id_fkey", "sessions", schema=schema) print("Dropped app_id foreign key for sessions table") - except Exception as e: - print(f"Error dropping app_id foreign key for sessions table: {e}") + else: + print("app_id foreign key for sessions table does not exist; skipping drop") # Drop the column op.drop_column("sessions", "app_id", schema=schema) diff --git a/migrations/versions/88b0fb10906f_add_webhooks_table.py b/migrations/versions/88b0fb10906f_add_webhooks_table.py index da607938..39b5bf51 100644 --- a/migrations/versions/88b0fb10906f_add_webhooks_table.py +++ b/migrations/versions/88b0fb10906f_add_webhooks_table.py @@ -183,3 +183,46 @@ def downgrade() -> None: if column_exists("active_queue_sessions", "work_unit_data", inspector): op.drop_column("active_queue_sessions", "work_unit_data", schema=schema) + + # Re-add columns to active_queue_sessions + if not column_exists("active_queue_sessions", "session_id", inspector): + op.add_column( + "active_queue_sessions", + sa.Column("session_id", sa.TEXT(), nullable=True), + schema=schema, + ) + + if not column_exists("active_queue_sessions", "sender_name", inspector): + op.add_column( + "active_queue_sessions", + sa.Column("sender_name", sa.TEXT(), nullable=True), + schema=schema, + ) + + if not column_exists("active_queue_sessions", "target_name", inspector): + op.add_column( + "active_queue_sessions", + sa.Column("target_name", sa.TEXT(), nullable=True), + schema=schema, + ) + + if not column_exists("active_queue_sessions", "task_type", inspector): + op.add_column( + "active_queue_sessions", + sa.Column("task_type", sa.TEXT(), nullable=False), + schema=schema, + ) + + # Re-add unique constraint + if not constraint_exists( + "active_queue_sessions", + "unique_active_queue_session", + "unique", + inspector, + ): + op.create_unique_constraint( + "unique_active_queue_session", + "active_queue_sessions", + ["session_id", "sender_name", "target_name", "task_type"], + schema=schema, + ) diff --git a/migrations/versions/917195d9b5e9_add_messageembedding_table.py b/migrations/versions/917195d9b5e9_add_messageembedding_table.py index 8d8f78f4..4eb73739 100644 --- a/migrations/versions/917195d9b5e9_add_messageembedding_table.py +++ b/migrations/versions/917195d9b5e9_add_messageembedding_table.py @@ -40,14 +40,15 @@ def upgrade() -> None: server_default=sa.func.now(), ), # Foreign key constraints - sa.ForeignKeyConstraint(["message_id"], ["messages.public_id"]), - sa.ForeignKeyConstraint(["workspace_name"], ["workspaces.name"]), + sa.ForeignKeyConstraint(["message_id"], [f"{schema}.messages.public_id"]), + sa.ForeignKeyConstraint(["workspace_name"], [f"{schema}.workspaces.name"]), sa.ForeignKeyConstraint( ["session_name", "workspace_name"], - ["sessions.name", "sessions.workspace_name"], + [f"{schema}.sessions.name", f"{schema}.sessions.workspace_name"], ), sa.ForeignKeyConstraint( - ["peer_name", "workspace_name"], ["peers.name", "peers.workspace_name"] + ["peer_name", "workspace_name"], + [f"{schema}.peers.name", f"{schema}.peers.workspace_name"], ), schema=schema, ) diff --git a/migrations/versions/a1b2c3d4e5f6_initial_schema.py b/migrations/versions/a1b2c3d4e5f6_initial_schema.py index dfb47843..cafa43ca 100644 --- a/migrations/versions/a1b2c3d4e5f6_initial_schema.py +++ b/migrations/versions/a1b2c3d4e5f6_initial_schema.py @@ -83,7 +83,7 @@ def upgrade() -> None: sa.CheckConstraint("length(name) <= 512", name="name_length"), sa.CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"), sa.ForeignKeyConstraint( - ["app_id"], ["apps.public_id"], name=op.f("fk_users_app_id_apps") + ["app_id"], [f"{schema}.apps.public_id"], name=op.f("fk_users_app_id_apps") ), sa.PrimaryKeyConstraint("id", name=op.f("pk_users")), sa.UniqueConstraint("name", "app_id", name="unique_name_app_user"), @@ -130,7 +130,9 @@ def upgrade() -> None: sa.CheckConstraint("length(public_id) = 21", name="public_id_length"), sa.CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"), sa.ForeignKeyConstraint( - ["user_id"], ["users.public_id"], name=op.f("fk_sessions_user_id_users") + ["user_id"], + [f"{schema}.users.public_id"], + name=op.f("fk_sessions_user_id_users"), ), sa.PrimaryKeyConstraint("id", name=op.f("pk_sessions")), sa.UniqueConstraint("public_id", name=op.f("uq_sessions_public_id")), @@ -186,7 +188,7 @@ def upgrade() -> None: sa.CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"), sa.ForeignKeyConstraint( ["session_id"], - ["sessions.public_id"], + [f"{schema}.sessions.public_id"], name=op.f("fk_messages_session_id_sessions"), ), sa.PrimaryKeyConstraint("id", name=op.f("pk_messages")), @@ -246,7 +248,7 @@ def upgrade() -> None: sa.CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"), sa.ForeignKeyConstraint( ["message_id"], - ["messages.public_id"], + [f"{schema}.messages.public_id"], name=op.f("fk_metamessages_message_id_messages"), ), sa.PrimaryKeyConstraint("id", name=op.f("pk_metamessages")), @@ -308,7 +310,9 @@ def upgrade() -> None: sa.CheckConstraint("length(name) <= 512", name="name_length"), sa.CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"), sa.ForeignKeyConstraint( - ["user_id"], ["users.public_id"], name=op.f("fk_collections_user_id_users") + ["user_id"], + [f"{schema}.users.public_id"], + name=op.f("fk_collections_user_id_users"), ), sa.PrimaryKeyConstraint("id", name=op.f("pk_collections")), sa.UniqueConstraint("name", "user_id", name="unique_name_collection_user"), @@ -372,7 +376,7 @@ def upgrade() -> None: sa.CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"), sa.ForeignKeyConstraint( ["collection_id"], - ["collections.public_id"], + [f"{schema}.collections.public_id"], name=op.f("fk_documents_collection_id_collections"), ), sa.PrimaryKeyConstraint("id", name=op.f("pk_documents")), @@ -412,7 +416,9 @@ def upgrade() -> None: sa.Column("payload", postgresql.JSONB(astext_type=sa.Text()), nullable=False), sa.Column("processed", sa.Boolean(), nullable=False, server_default="false"), sa.ForeignKeyConstraint( - ["session_id"], ["sessions.id"], name=op.f("fk_queue_session_id_sessions") + ["session_id"], + [f"{schema}.sessions.id"], + name=op.f("fk_queue_session_id_sessions"), ), sa.PrimaryKeyConstraint("id", name=op.f("pk_queue")), schema=schema, @@ -437,7 +443,7 @@ def upgrade() -> None: ), sa.ForeignKeyConstraint( ["session_id"], - ["sessions.id"], + [f"{schema}.sessions.id"], name=op.f("fk_active_queue_sessions_session_id_sessions"), ), sa.PrimaryKeyConstraint("session_id", name=op.f("pk_active_queue_sessions")), diff --git a/migrations/versions/d429de0e5338_adopt_peer_paradigm.py b/migrations/versions/d429de0e5338_adopt_peer_paradigm.py index 9a8d9844..6cb42b74 100644 --- a/migrations/versions/d429de0e5338_adopt_peer_paradigm.py +++ b/migrations/versions/d429de0e5338_adopt_peer_paradigm.py @@ -13,7 +13,7 @@ import sqlalchemy as sa import tiktoken from alembic import op from nanoid import generate as generate_nanoid -from sqlalchemy import text +from sqlalchemy import Inspector, text from sqlalchemy.dialects import postgresql from migrations.utils import ( @@ -114,8 +114,117 @@ def downgrade() -> None: # Step 11: Restore foreign keys restore_foreign_keys(schema) + # Step 12: Readd metamessages table + if not table_exists("metamessages", inspector): + op.create_table( + "metamessages", + sa.Column( + "id", + sa.BigInteger(), + sa.Identity(always=False), + primary_key=True, + autoincrement=True, + index=True, + nullable=False, + ), + sa.Column("public_id", sa.TEXT(), nullable=False, index=True, unique=True), + sa.Column("content", sa.TEXT(), nullable=True), + sa.Column( + "message_id", + sa.TEXT(), + sa.ForeignKey("messages.public_id"), + nullable=True, + index=True, + ), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + nullable=False, + index=True, + server_default=sa.text("now()"), + ), + sa.Column("label", sa.TEXT(), nullable=False, index=True), + sa.Column( + "session_id", + sa.TEXT(), + sa.ForeignKey("sessions.public_id"), + nullable=True, + index=True, + ), + sa.Column( + "user_id", + sa.TEXT(), + sa.ForeignKey("users.public_id"), + nullable=True, + index=True, + ), + sa.Column( + "app_id", + sa.TEXT(), + sa.ForeignKey("apps.public_id"), + nullable=True, + index=True, + ), + sa.Column( + "metadata", + postgresql.JSONB(astext_type=sa.Text()), + nullable=False, + server_default=sa.text("'{}'::jsonb"), + ), + sa.Index( + "idx_metamessages_lookup", + "label", + sa.text("id DESC"), + postgresql_include=["public_id", "message_id", "created_at"], + ), + sa.Index( + "idx_metamessages_user_lookup", + "user_id", + "label", + sa.text("id DESC"), + ), + sa.Index( + "idx_metamessages_session_lookup", + "session_id", + "label", + sa.text("id DESC"), + ), + sa.Index( + "idx_metamessages_message_lookup", + "message_id", + "label", + sa.text("id DESC"), + ), + sa.CheckConstraint("length(public_id) = 21", name="public_id_length"), + sa.CheckConstraint( + "public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format" + ), + sa.CheckConstraint("length(content) <= 65535", name="content_length"), + sa.CheckConstraint("length(label) <= 512", name="label_length"), + # Added constraints to ensure consistency + sa.CheckConstraint( + "(message_id IS NULL) OR (session_id IS NOT NULL)", + name="message_requires_session", + ), + sa.ForeignKeyConstraint( + ["session_id"], + [f"{schema}.sessions.public_id"], + "fk_metamessages_session_id_sessions", + ), + sa.ForeignKeyConstraint( + ["user_id"], + [f"{schema}.users.public_id"], + "fk_metamessages_user_id_users", + ), + sa.ForeignKeyConstraint( + ["app_id"], + [f"{schema}.apps.public_id"], + "fk_metamessages_app_id_apps", + ), + ) -def rename_tables(schema: str, inspector) -> None: + +def rename_tables(schema: str, inspector: Inspector) -> None: """Rename apps->workspaces and users->peers tables.""" if inspector.has_table("apps", schema=schema): op.rename_table("apps", "workspaces", schema=schema) @@ -123,7 +232,7 @@ def rename_tables(schema: str, inspector) -> None: op.rename_table("users", "peers", schema=schema) -def update_workspaces_table(schema: str, inspector) -> None: +def update_workspaces_table(schema: str, inspector: Inspector) -> None: """Update workspaces table (formerly apps).""" # Add configuration column @@ -176,7 +285,7 @@ def update_workspaces_table(schema: str, inspector) -> None: ) -def update_peers_table(schema: str, inspector) -> None: +def update_peers_table(schema: str, inspector: Inspector) -> None: """Update peers table (formerly users).""" # Add configuration column @@ -266,9 +375,10 @@ def update_peers_table(schema: str, inspector) -> None: ) op.drop_column("peers", "app_id", schema=schema) + op.alter_column("peers", "workspace_name", nullable=False, schema=schema) -def update_sessions_table(schema: str, inspector) -> None: +def update_sessions_table(schema: str, inspector: Inspector) -> None: """Update sessions table.""" # Add configuration column @@ -304,6 +414,7 @@ def update_sessions_table(schema: str, inspector) -> None: f"UPDATE {schema}.sessions SET workspace_name = workspaces.name FROM {schema}.workspaces WHERE sessions.app_id = workspaces.id" ) ) + op.alter_column("sessions", "workspace_name", nullable=False, schema=schema) op.add_column( "sessions", @@ -365,7 +476,7 @@ def update_sessions_table(schema: str, inspector) -> None: ) -def create_and_populate_session_peers_table(schema: str, inspector) -> None: +def create_and_populate_session_peers_table(schema: str, inspector: Inspector) -> None: """Create and populate session_peers table.""" # Create session_peers table @@ -400,11 +511,11 @@ def create_and_populate_session_peers_table(schema: str, inspector) -> None: ), sa.ForeignKeyConstraint( ["peer_name", "workspace_name"], - ["peers.name", "peers.workspace_name"], + [f"{schema}.peers.name", f"{schema}.peers.workspace_name"], ), sa.ForeignKeyConstraint( ["session_name", "workspace_name"], - ["sessions.name", "sessions.workspace_name"], + [f"{schema}.sessions.name", f"{schema}.sessions.workspace_name"], ), sa.PrimaryKeyConstraint("workspace_name", "session_name", "peer_name"), ) @@ -456,7 +567,7 @@ def create_and_populate_session_peers_table(schema: str, inspector) -> None: ) -def update_messages_table(schema: str, inspector) -> None: +def update_messages_table(schema: str, inspector: Inspector) -> None: """Update messages table.""" # Add new columns @@ -618,7 +729,7 @@ def update_messages_table(schema: str, inspector) -> None: backfill_token_counts(schema) -def update_collections_table(schema: str, inspector) -> None: +def update_collections_table(schema: str, inspector: Inspector) -> None: """Update collections table.""" # Add new columns @@ -741,7 +852,7 @@ def update_collections_table(schema: str, inspector) -> None: ) -def update_documents_table(schema: str, inspector) -> None: +def update_documents_table(schema: str, inspector: Inspector) -> None: """Update documents table.""" # Add new columns @@ -873,7 +984,9 @@ def update_documents_table(schema: str, inspector) -> None: ) -def update_queue_and_active_queue_sessions_tables(schema: str, inspector) -> None: +def update_queue_and_active_queue_sessions_tables( + schema: str, inspector: Inspector +) -> None: """Update queue and active_queue_sessions tables.""" # Drop foreign key constraints first, before changing column types @@ -903,7 +1016,7 @@ def update_queue_and_active_queue_sessions_tables(schema: str, inspector) -> Non # Get the mapping of old session.id (integer) to new session.id (text, which is public_id) # At this point, sessions table still has both id (integer) and public_id (text) columns - session_id_mapping = {} + session_id_mapping: dict[int, str] = {} if table_exists("sessions", inspector): sessions_mapping = connection.execute( sa.text(f"SELECT id, public_id FROM {schema}.sessions") @@ -1081,7 +1194,7 @@ def backfill_token_counts(schema: str) -> None: break # Calculate token counts and update messages in batches - batch_updates = [] + batch_updates: list[tuple[int, int]] = [] for message_id, content in messages: token_count = _count_tokens(content) batch_updates.append((message_id, token_count)) @@ -1110,7 +1223,7 @@ def backfill_token_counts(schema: str) -> None: offset += batch_size -def restore_app_user_columns(schema: str, inspector) -> None: +def restore_app_user_columns(schema: str, inspector: Inspector) -> None: """Restore app_id and user_id columns to peers and sessions.""" # Add app_id back to peers if not column_exists("peers", "app_id", inspector): @@ -1154,8 +1267,7 @@ def restore_app_user_columns(schema: str, inspector) -> None: schema=schema, ) # Populate user_id from session_peers - # The user peer is the one that existed before agent peers were created - # Agent peers have names that match their IDs (both are nanoids), user peers have different names and IDs + # Prefer user peers (where id != name), but fall back to any peer if needed op.execute( sa.text(f""" UPDATE {schema}.sessions SET user_id = ( @@ -1172,7 +1284,7 @@ def restore_app_user_columns(schema: str, inspector) -> None: op.alter_column("sessions", "user_id", nullable=False, schema=schema) -def restore_documents_table(schema: str, inspector) -> None: +def restore_documents_table(schema: str, inspector: Inspector) -> None: """Restore documents table to pre-peer paradigm state.""" # Add back id column as primary key if not column_exists("documents", "temp_id", inspector): @@ -1306,7 +1418,7 @@ def restore_documents_table(schema: str, inspector) -> None: ) -def restore_collections_table(schema: str, inspector) -> None: +def restore_collections_table(schema: str, inspector: Inspector) -> None: """Restore collections table to pre-peer paradigm state.""" # Add back id column as primary key if not column_exists("collections", "temp_id", inspector): @@ -1433,7 +1545,7 @@ def restore_collections_table(schema: str, inspector) -> None: ) -def restore_messages_table(schema: str, inspector) -> None: +def restore_messages_table(schema: str, inspector: Inspector) -> None: """Restore messages table to pre-peer paradigm state.""" # Add back old columns if not column_exists("messages", "session_id", inspector): @@ -1553,8 +1665,20 @@ def restore_messages_table(schema: str, inspector) -> None: referent_schema=schema, ) + op.create_index( + "idx_messages_session_lookup", + "messages", + ["session_id", "id"], + postgresql_include=[ + "public_id", + "is_user", + "created_at", + ], + schema=schema, + ) -def restore_sessions_table(schema: str, inspector) -> None: + +def restore_sessions_table(schema: str, inspector: Inspector) -> None: """Restore sessions table to pre-peer paradigm state.""" # Add back id column as BigInteger primary key if not column_exists("sessions", "temp_id", inspector): @@ -1624,8 +1748,12 @@ def restore_sessions_table(schema: str, inspector) -> None: "public_id_format", "sessions", "public_id ~ '^[A-Za-z0-9_-]+$'", schema=schema ) + op.create_index( + "idx_sessions_user_lookup", "sessions", ["user_id", "public_id"], schema=schema + ) -def restore_peers_table(schema: str, inspector) -> None: + +def restore_peers_table(schema: str, inspector: Inspector) -> None: """Restore peers table to pre-user paradigm state.""" # Add back id column as BigInteger primary key if not column_exists("peers", "temp_id", inspector): @@ -1703,7 +1831,7 @@ def restore_peers_table(schema: str, inspector) -> None: ) -def restore_workspaces_table(schema: str, inspector) -> None: +def restore_workspaces_table(schema: str, inspector: Inspector) -> None: """Restore workspaces table to pre-peer paradigm state (apps).""" # Add back id column as BigInteger primary key if not column_exists("workspaces", "temp_id", inspector): @@ -1764,14 +1892,16 @@ def restore_workspaces_table(schema: str, inspector) -> None: ) -def restore_queue_and_active_queue_sessions_tables(schema: str, inspector) -> None: +def restore_queue_and_active_queue_sessions_tables( + schema: str, inspector: Inspector +) -> None: """Restore queue and active_queue_sessions tables to pre-peer paradigm state.""" connection = op.get_bind() # Create reverse mapping from session.public_id (text) back to session.id (BigInteger) # At this point in downgrade, sessions table still has both id (BigInteger) and public_id (text) - session_id_reverse_mapping = {} + session_id_reverse_mapping: dict[str, int] = {} if table_exists("sessions", inspector): sessions_mapping = connection.execute( sa.text(f"SELECT id, public_id FROM {schema}.sessions") @@ -1781,26 +1911,27 @@ def restore_queue_and_active_queue_sessions_tables(schema: str, inspector) -> No session_id_reverse_mapping[text_id] = big_int_id # Update queue table - if table_exists("queue", inspector) and session_id_reverse_mapping: + if table_exists("queue", inspector): # Get current session_id values in queue table (they are text now) - queue_session_ids = connection.execute( - sa.text( - f"SELECT DISTINCT session_id FROM {schema}.queue WHERE session_id IS NOT NULL" - ) - ).fetchall() - - # Convert session_id values back to BigInteger - for (session_id,) in queue_session_ids: - if session_id in session_id_reverse_mapping: - old_id = session_id_reverse_mapping[session_id] - connection.execute( - sa.text( - f"UPDATE {schema}.queue SET session_id = :old_id WHERE session_id = :new_id" - ), - {"old_id": str(old_id), "new_id": str(session_id)}, + if session_id_reverse_mapping: + queue_session_ids = connection.execute( + sa.text( + f"SELECT DISTINCT session_id FROM {schema}.queue WHERE session_id IS NOT NULL" ) + ).fetchall() - # Change column type back to BigInteger + # Convert session_id values back to BigInteger + for (session_id,) in queue_session_ids: + if session_id in session_id_reverse_mapping: + old_id = session_id_reverse_mapping[session_id] + connection.execute( + sa.text( + f"UPDATE {schema}.queue SET session_id = :old_id WHERE session_id = :new_id" + ), + {"old_id": str(old_id), "new_id": str(session_id)}, + ) + + # Change column type back to BigInteger (always) op.alter_column( "queue", "session_id", @@ -1810,36 +1941,40 @@ def restore_queue_and_active_queue_sessions_tables(schema: str, inspector) -> No ) # Update active_queue_sessions table - if table_exists("active_queue_sessions", inspector) and session_id_reverse_mapping: - # Get current session_id values in active_queue_sessions table (they are text now) - active_queue_session_ids = connection.execute( - sa.text( - f"SELECT DISTINCT session_id FROM {schema}.active_queue_sessions WHERE session_id IS NOT NULL" - ) - ).fetchall() + if table_exists("active_queue_sessions", inspector): + if session_id_reverse_mapping: + # Get current session_id values in active_queue_sessions table (they are text now) + active_queue_session_ids = connection.execute( + sa.text( + f"SELECT DISTINCT session_id FROM {schema}.active_queue_sessions WHERE session_id IS NOT NULL" + ) + ).fetchall() - if constraint_exists( - "active_queue_sessions", "pk_active_queue_sessions", "primary", inspector - ): - op.drop_constraint( - "pk_active_queue_sessions", + if constraint_exists( "active_queue_sessions", - type_="primary", - schema=schema, - ) - - # Convert session_id values back to BigInteger - for (session_id,) in active_queue_session_ids: - if session_id in session_id_reverse_mapping: - old_id = session_id_reverse_mapping[session_id] - connection.execute( - sa.text( - f"UPDATE {schema}.active_queue_sessions SET session_id = :old_id WHERE session_id = :new_id" - ), - {"old_id": str(old_id), "new_id": str(session_id)}, + "pk_active_queue_sessions", + "primary", + inspector, + ): + op.drop_constraint( + "pk_active_queue_sessions", + "active_queue_sessions", + type_="primary", + schema=schema, ) - # Change column type back to BigInteger + # Convert session_id values back to BigInteger + for (session_id,) in active_queue_session_ids: + if session_id in session_id_reverse_mapping: + old_id = session_id_reverse_mapping[session_id] + connection.execute( + sa.text( + f"UPDATE {schema}.active_queue_sessions SET session_id = :old_id WHERE session_id = :new_id" + ), + {"old_id": str(old_id), "new_id": str(session_id)}, + ) + + # Change column type back to BigInteger (always) op.alter_column( "active_queue_sessions", "session_id", @@ -1848,12 +1983,16 @@ def restore_queue_and_active_queue_sessions_tables(schema: str, inspector) -> No postgresql_using="session_id::bigint", ) - op.create_primary_key( - "pk_active_queue_sessions", - "active_queue_sessions", - ["session_id"], - schema=schema, - ) + # Ensure primary key exists on session_id for pre-peer shape + if not constraint_exists( + "active_queue_sessions", "pk_active_queue_sessions", "primary", inspector + ): + op.create_primary_key( + "pk_active_queue_sessions", + "active_queue_sessions", + ["session_id"], + schema=schema, + ) if constraint_exists( "active_queue_sessions", @@ -1877,29 +2016,10 @@ def restore_queue_and_active_queue_sessions_tables(schema: str, inspector) -> No if column_exists("active_queue_sessions", "task_type", inspector): op.drop_column("active_queue_sessions", "task_type", schema=schema) - # Restore foreign key constraints - if table_exists("queue", inspector): - op.create_foreign_key( - "fk_queue_session_id_sessions", - "queue", - "sessions", - ["session_id"], - ["id"], - referent_schema=schema, - ) - - if table_exists("active_queue_sessions", inspector): - op.create_foreign_key( - "fk_active_queue_sessions_session_id_sessions", - "active_queue_sessions", - "sessions", - ["session_id"], - ["id"], - referent_schema=schema, - ) + # Defer restoring foreign key constraints to restore_foreign_keys() -def restore_table_names(schema: str, inspector) -> None: +def restore_table_names(schema: str, inspector: Inspector) -> None: """Restore table names: workspaces->apps and peers->users.""" if inspector.has_table("workspaces", schema=schema): op.rename_table("workspaces", "apps", schema=schema) @@ -1957,6 +2077,23 @@ def restore_foreign_keys(schema: str) -> None: ["public_id"], referent_schema=schema, ) + # Restore queue/session FKs last, after all type changes are complete + op.create_foreign_key( + "fk_queue_session_id_sessions", + "queue", + "sessions", + ["session_id"], + ["id"], + referent_schema=schema, + ) + op.create_foreign_key( + "fk_active_queue_sessions_session_id_sessions", + "active_queue_sessions", + "sessions", + ["session_id"], + ["id"], + referent_schema=schema, + ) op.create_foreign_key( "fk_sessions_app_id_apps", "sessions", diff --git a/pyproject.toml b/pyproject.toml index 9b0de576..0022cd1e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -81,7 +81,7 @@ extend-immutable-calls = ["fastapi.Depends"] [tool.pytest.ini_options] asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "session" -addopts = "--strict-markers --cov=src/ --cov=sdks/python/src/honcho --cov-report=term-missing" +addopts = "--strict-markers --cov=src/ --cov=sdks/python/src/honcho --cov-report=term-missing --ignore=tests/alembic" testpaths = ["tests"] pythonpath = ["src"] diff --git a/scripts/ensure_alembic_tests.py b/scripts/ensure_alembic_tests.py new file mode 100644 index 00000000..fe3cad8b --- /dev/null +++ b/scripts/ensure_alembic_tests.py @@ -0,0 +1,68 @@ +#!/usr/bin/env uv run python +""" +Script that validates that all alembic migration revisions have a corresponding test file. +Note that this script is actively used within our precommit hooks and should not be removed. +If this script is moved, the corresponding precommit hook will need to be updated. +""" + +import sys +from pathlib import Path + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +MIGRATIONS_DIR = PROJECT_ROOT / "migrations" / "versions" +TESTS_DIR = PROJECT_ROOT / "tests" / "alembic" / "revisions" + + +def main() -> int: + missing: list[str] = [] + + # Gather migration base names (without extension) + migration_files = [ + p for p in MIGRATIONS_DIR.glob("*.py") if p.name != "__init__.py" + ] + migration_basenames = {p.stem for p in migration_files} + print(f"Migration basenames: {migration_basenames}") + + # Gather test file base names mapped back to migration names by stripping leading 'test_' + test_files = [p for p in TESTS_DIR.glob("test_*.py") if p.name != "__init__.py"] + test_targets = {p.stem.removeprefix("test_") for p in test_files} + + for migration_basename in sorted(migration_basenames): + if migration_basename not in test_targets: + missing.append(migration_basename) + if missing: + print( + "Missing Alembic tests for the following migration revisions:", + file=sys.stderr, + ) + for name in missing: + print(f" - {name}", file=sys.stderr) + print( + "\nExpected test files under tests/alembic/revisions named as:", + file=sys.stderr, + ) + for name in missing: + print(f" - tests/alembic/revisions/test_{name}.py", file=sys.stderr) + print("\nScaffold helper commands:", file=sys.stderr) + for name in missing: + print( + f" - python -m tests.alembic.scaffold {name.split('_', 1)[0]}", + file=sys.stderr, + ) + return 1 + + # Optional: warn if there are tests without corresponding migrations (stale tests) + stale_tests = sorted(test_targets - migration_basenames) + if stale_tests: + print( + "Warning: Found tests without corresponding migration files (stale?):", + file=sys.stderr, + ) + for name in stale_tests: + print(f" - tests/alembic/revisions/test_{name}.py", file=sys.stderr) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/generate_jwt_secret.py b/scripts/generate_jwt_secret.py index 449d4cc7..6a7dc5a4 100755 --- a/scripts/generate_jwt_secret.py +++ b/scripts/generate_jwt_secret.py @@ -1,4 +1,4 @@ -#!/usr/bin/env python +#!/usr/bin/env uv run python """ Utility script to generate a JWT secret for use in the .env file. This uses the same logic as the automatically generated version in security.py. diff --git a/scripts/generate_message_embeddings.py b/scripts/generate_message_embeddings.py index fc19c3fe..75c3f22f 100644 --- a/scripts/generate_message_embeddings.py +++ b/scripts/generate_message_embeddings.py @@ -1,3 +1,4 @@ +#!/usr/bin/env uv run python """ Script to generate embeddings for existing messages that don't already have embeddings. diff --git a/src/models.py b/src/models.py index 1cd3b00c..b66e31cf 100644 --- a/src/models.py +++ b/src/models.py @@ -113,7 +113,7 @@ class Peer(Base): DateTime(timezone=True), index=True, default=func.now() ) workspace_name: Mapped[str] = mapped_column( - ForeignKey("workspaces.name"), index=True + ForeignKey("workspaces.name"), index=True, nullable=False ) configuration: Mapped[dict[str, Any]] = mapped_column(JSONB, default=dict) @@ -149,7 +149,7 @@ class Session(Base): ) messages = relationship("Message", back_populates="session") workspace_name: Mapped[str] = mapped_column( - ForeignKey("workspaces.name"), index=True + ForeignKey("workspaces.name"), index=True, nullable=False ) configuration: Mapped[dict[str, Any]] = mapped_column(JSONB, default=dict) diff --git a/tests/alembic/README.md b/tests/alembic/README.md new file mode 100644 index 00000000..2dec3d36 --- /dev/null +++ b/tests/alembic/README.md @@ -0,0 +1,29 @@ +These tests validate Alembic migrations end-to-end for structure, order, and correctness. +They ensure reversibility, expected schema and data, and integration with the registry and pipeline. +The key components are the verifier, the test pipeline, the registry, and the revisions under test. + +### Verifier + +- The verifier runs checks when specific revisions are applied and reverted. +- It validates the schema after upgrade, verifies data migrations such as defaults, backfills, and transforms, and confirms reversibility after downgrade. +- Assertions are grouped per-revision or feature, and helpers use the SQLAlchemy inspector to introspect the database. + +### Test Pipeline + +- The pipeline orchestrates the database lifecycle by creating a database, applying upgrades and downgrades, running verifications, and tearing down resources. +- It typically starts from base, upgrades to the revision immediately before a target revision, seeds the DB, runs the target migration, and then validates the schema + data +- It relies on shared fixtures such as `engine`, `connection`, and `alembic_config`, and it ensures isolation per test + +### Registry + +- The registry declares revisions and test metadata used to drive scenarios. +- It defines ordering and selection, attaches verifier callbacks to specific revisions or ranges, and centralizes per revision expectations. + +### Revisions + +- Revisions are the migration scripts under `alembic/revisions`. +- Each revision should provide functions decorated with `register_before_upgrade()` and `register_after_upgrade()`. These are used to validate schemas and data before and after a migration is run + +### Running the Tests + +- Tests can be run all together `pytest tests/alembic` or individually `pytest tests/alembic -k "revision_number"`. For example, to run the test against a1b2c3d4e5f6_initial_schema.py, you would run the command `pytest tests/alembic -k "a1b2c3d4e5f6"` diff --git a/tests/alembic/__init__.py b/tests/alembic/__init__.py new file mode 100644 index 00000000..571f7c81 --- /dev/null +++ b/tests/alembic/__init__.py @@ -0,0 +1,5 @@ +"""Test package for Alembic migration scenarios.""" + +from . import revisions + +__all__ = ["revisions"] diff --git a/tests/alembic/conftest.py b/tests/alembic/conftest.py new file mode 100644 index 00000000..022ae5fc --- /dev/null +++ b/tests/alembic/conftest.py @@ -0,0 +1,107 @@ +"""Pytest configuration for Alembic-focused tests.""" + +from __future__ import annotations + +import os +from collections.abc import Generator +from pathlib import Path + +import pytest +from alembic.config import Config +from sqlalchemy import create_engine, text +from sqlalchemy.engine import Engine +from sqlalchemy.engine.url import URL +from sqlalchemy_utils import ( + create_database, # pyright: ignore + database_exists, # pyright: ignore + drop_database, # pyright: ignore +) + +from src.config import settings +from tests.conftest import CONNECTION_URI + +ALEMBIC_CONFIG_PATH = Path(__file__).resolve().parents[2] / "alembic.ini" +ALEMBIC_TEST_DB_URL: URL = CONNECTION_URI.set(database="alembic_migration_tests") + + +@pytest.fixture(scope="session", autouse=True) +def configure_alembic_settings(alembic_database: str) -> Generator[None, None, None]: + """Point application settings at the Alembic test database.""" + + previous_uri = settings.DB.CONNECTION_URI + os.environ["DB_CONNECTION_URI"] = alembic_database + settings.DB.CONNECTION_URI = alembic_database + + try: + yield + finally: + settings.DB.CONNECTION_URI = previous_uri + if previous_uri: + os.environ["DB_CONNECTION_URI"] = previous_uri + else: + os.environ.pop("DB_CONNECTION_URI", None) + + +@pytest.fixture +def alembic_cfg(alembic_database: str) -> Config: + """Provide an Alembic Config bound to the alembic test database.""" + cfg = Config(str(ALEMBIC_CONFIG_PATH)) + cfg.set_main_option( + "script_location", str(ALEMBIC_CONFIG_PATH.parent / "migrations") + ) + cfg.set_main_option("sqlalchemy.url", alembic_database) + return cfg + + +@pytest.fixture(scope="session") +def alembic_database() -> Generator[str, None, None]: + """Provision a dedicated DB for Alembic verification tests.""" + + assert ALEMBIC_TEST_DB_URL.database == "alembic_migration_tests", ( + "Can't set up Alembic test database fixture. " + + "ALEMBIC_TEST_DB_URL.database is {ALEMBIC_TEST_DB_URL.database}, " + + "expected 'alembic_migration_tests'. " + ) + + if database_exists(ALEMBIC_TEST_DB_URL): + drop_database(ALEMBIC_TEST_DB_URL) # start fresh + create_database(ALEMBIC_TEST_DB_URL) + + engine = create_engine(str(ALEMBIC_TEST_DB_URL)) + try: + with engine.begin() as conn: + conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + yield str(ALEMBIC_TEST_DB_URL) + finally: + engine.dispose() + if database_exists(ALEMBIC_TEST_DB_URL): + drop_database(ALEMBIC_TEST_DB_URL) + + +@pytest.fixture +def alembic_engine(alembic_database: str) -> Generator[Engine, None, None]: + """Yield an engine bound to the Alembic test database.""" + + engine = create_engine(alembic_database, pool_pre_ping=True) + try: + yield engine + finally: + engine.dispose() + + +@pytest.fixture(autouse=True) +def reset_schema_between_tests( + alembic_engine: Engine, +) -> Generator[None, None, None]: + """Drop and recreate the schema for a clean slate each test (fast reset).""" + + schema = settings.DB.SCHEMA + + def _reset_schema() -> None: + with alembic_engine.begin() as conn: + conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE')) + conn.execute(text(f'CREATE SCHEMA "{schema}"')) + + _reset_schema() + yield + _reset_schema() diff --git a/tests/alembic/registry.py b/tests/alembic/registry.py new file mode 100644 index 00000000..d2b013bc --- /dev/null +++ b/tests/alembic/registry.py @@ -0,0 +1,73 @@ +"""Registry for migration-specific prepare and verification hooks.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from tests.alembic.verifier import MigrationVerifier + + +BeforeUpgradeHook = Callable[["MigrationVerifier"], None] +AfterUpgradeHook = Callable[["MigrationVerifier"], None] + + +@dataclass(slots=True) +class RevisionHooks: + """Container for lifecycle hooks tied to a revision.""" + + before_upgrade: BeforeUpgradeHook | None = None + after_upgrade: AfterUpgradeHook | None = None + + +_REGISTRY: dict[str, RevisionHooks] = {} + + +def register_before_upgrade( + revision: str, +) -> Callable[[BeforeUpgradeHook], BeforeUpgradeHook]: + """Register a callable to execute before upgrading to ``revision``.""" + + def decorator(func: BeforeUpgradeHook) -> BeforeUpgradeHook: + hooks = _REGISTRY.setdefault(revision, RevisionHooks()) + if hooks.before_upgrade is not None: + raise ValueError( + f"before_upgrade hook already registered for revision {revision}" + ) + hooks.before_upgrade = func + return func + + return decorator + + +def register_after_upgrade( + revision: str, +) -> Callable[[AfterUpgradeHook], AfterUpgradeHook]: + """Register a callable to execute after upgrading to ``revision``.""" + + def decorator(func: AfterUpgradeHook) -> AfterUpgradeHook: + hooks = _REGISTRY.setdefault(revision, RevisionHooks()) + if hooks.after_upgrade is not None: + raise ValueError( + f"after_upgrade hook already registered for revision {revision}" + ) + hooks.after_upgrade = func + return func + + return decorator + + +def get_registered_hooks() -> Mapping[str, RevisionHooks]: + """Return a snapshot of the registered revision hooks.""" + + return dict(_REGISTRY) + + +__all__ = [ + "RevisionHooks", + "register_before_upgrade", + "register_after_upgrade", + "get_registered_hooks", +] diff --git a/tests/alembic/revisions/__init__.py b/tests/alembic/revisions/__init__.py new file mode 100644 index 00000000..359e7b1f --- /dev/null +++ b/tests/alembic/revisions/__init__.py @@ -0,0 +1,35 @@ +"""Register revision-specific hooks for migration verification.""" + +from . import ( + test_05486ce795d5_make_session_name_required_on_messages, + test_08894082221a_replace_collection_name_with_observer_, + test_20f89a421aff_rename_metamessage_type_to_label, + test_66e63cf2cf77_add_indexes_to_documents_table, + test_76ffba56fe8c_add_error_field_to_queueitem, + test_88b0fb10906f_add_webhooks_table, + test_556a16564f50_add_user_id_and_app_id_to_tables, + test_564ba40505c5_add_session_name_column_to_documents, + test_917195d9b5e9_add_messageembedding_table, + test_a1b2c3d4e5f6_initial_schema, + test_b765d82110bd_change_metamessages_to_user_level_with_, + test_bb6fb3a7a643_add_message_seq_in_session_column, + test_c3828084f472_add_indexes_for_messages_and_, + test_d429de0e5338_adopt_peer_paradigm, +) + +__all__ = [ + "test_05486ce795d5_make_session_name_required_on_messages", + "test_08894082221a_replace_collection_name_with_observer_", + "test_20f89a421aff_rename_metamessage_type_to_label", + "test_556a16564f50_add_user_id_and_app_id_to_tables", + "test_564ba40505c5_add_session_name_column_to_documents", + "test_66e63cf2cf77_add_indexes_to_documents_table", + "test_76ffba56fe8c_add_error_field_to_queueitem", + "test_88b0fb10906f_add_webhooks_table", + "test_917195d9b5e9_add_messageembedding_table", + "test_a1b2c3d4e5f6_initial_schema", + "test_b765d82110bd_change_metamessages_to_user_level_with_", + "test_bb6fb3a7a643_add_message_seq_in_session_column", + "test_c3828084f472_add_indexes_for_messages_and_", + "test_d429de0e5338_adopt_peer_paradigm", +] diff --git a/tests/alembic/revisions/test_05486ce795d5_make_session_name_required_on_messages.py b/tests/alembic/revisions/test_05486ce795d5_make_session_name_required_on_messages.py new file mode 100644 index 00000000..093ed17b --- /dev/null +++ b/tests/alembic/revisions/test_05486ce795d5_make_session_name_required_on_messages.py @@ -0,0 +1,104 @@ +"""Hooks for revision 05486ce795d5 (session_name required).""" + +from __future__ import annotations + +from nanoid import generate as generate_nanoid +from sqlalchemy import text + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + +WORKSPACE_ID = generate_nanoid() +WORKSPACE_NAME = "workspace-name" +PEER_ID = generate_nanoid() +PEER_NAME = "peer-name" +SESSION_ID = generate_nanoid() +SESSION_NAME = "session-name" +MESSAGE_ID = generate_nanoid() +MESSAGE_CONTENT = "message-content" + + +@register_before_upgrade("05486ce795d5") +def seed_orphaned_message(verifier: MigrationVerifier) -> None: + """Insert a message without a session_name to exercise the migration.""" + + schema = verifier.schema + connection = verifier.conn + + connection.execute( + text( + f""" + INSERT INTO "{schema}"."workspaces" ("id", "name") + VALUES (:workspace_id, :workspace_name) + """ + ), + {"workspace_id": WORKSPACE_ID, "workspace_name": WORKSPACE_NAME}, + ) + + connection.execute( + text( + f""" + INSERT INTO "{schema}"."peers" ("id", "name", "workspace_name") + VALUES (:peer_id, :peer_name, :workspace_name) + """ + ), + { + "peer_id": PEER_ID, + "peer_name": PEER_NAME, + "workspace_name": WORKSPACE_NAME, + }, + ) + + connection.execute( + text( + f""" + INSERT INTO "{schema}"."messages" + ("public_id", "workspace_name", "peer_name", "session_name", "content") + VALUES (:public_id, :workspace_name, :peer_name, NULL, :content) + """ + ), + { + "public_id": MESSAGE_ID, + "workspace_name": WORKSPACE_NAME, + "peer_name": PEER_NAME, + "content": MESSAGE_CONTENT, + }, + ) + + +@register_after_upgrade("05486ce795d5") +def verify_session_name_enforced(verifier: MigrationVerifier) -> None: + """Ensure session_name is populated and constrained after the migration.""" + + verifier.assert_column_exists("messages", "session_name", nullable=False) + + schema = verifier.schema + conn = verifier.conn + + session = conn.execute( + text( + f'SELECT "name", "workspace_name" FROM "{schema}"."sessions" ' + + 'WHERE "name" = :session_name AND "workspace_name" = :workspace_name' + ), + { + "session_name": PEER_NAME, + "workspace_name": WORKSPACE_NAME, + }, + ).one_or_none() + assert ( + session is not None + ), "Session with expected name and workspace does not exist" + + message = conn.execute( + text( + f'SELECT "public_id", "content" FROM "{schema}"."messages" ' + + 'WHERE "session_name" = :session_name AND "workspace_name" = :workspace_name AND "peer_name" = :peer_name' + ), + { + "session_name": PEER_NAME, + "workspace_name": WORKSPACE_NAME, + "peer_name": PEER_NAME, + }, + ).one_or_none() + assert message is not None, "Message with expected session_name does not exist" + assert message.content == MESSAGE_CONTENT diff --git a/tests/alembic/revisions/test_08894082221a_replace_collection_name_with_observer_.py b/tests/alembic/revisions/test_08894082221a_replace_collection_name_with_observer_.py new file mode 100644 index 00000000..30f9aeb6 --- /dev/null +++ b/tests/alembic/revisions/test_08894082221a_replace_collection_name_with_observer_.py @@ -0,0 +1,361 @@ +"""Hooks for revision 08894082221a (observer/observed refactor).""" + +from __future__ import annotations + +from nanoid import generate as generate_nanoid +from sqlalchemy import inspect, text + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + +COLLECTION_ID: str = generate_nanoid() +GLOBAL_REP_COLLECTION_NAME = "global_representation" +DOCUMENT_ID = generate_nanoid() +DOCUMENT_NAME = "document-name" +LEGACY_PEER_ID = generate_nanoid() +LEGACY_PEER_NAME = "legacy-peer" +WORKSPACE_ID = generate_nanoid() +WORKSPACE_NAME = "workspace-name" + + +# Based on migration logic: when name='global_representation', observer=peer_name and observed=peer_name +EXPECTED_OBSERVED = LEGACY_PEER_NAME + + +@register_before_upgrade("08894082221a") +def prepare_observer_observed(verifier: MigrationVerifier) -> None: + schema = verifier.schema + conn = verifier.conn + + # Create workspace + conn.execute( + text( + f'INSERT INTO "{schema}"."workspaces" ("id", "name") ' + + "VALUES (:workspace_id, :workspace_name)" + ), + {"workspace_id": WORKSPACE_ID, "workspace_name": WORKSPACE_NAME}, + ) + + # Create peer + conn.execute( + text( + f'INSERT INTO "{schema}"."peers" ("id", "name", "workspace_name") ' + + "VALUES (:peer_id, :peer_name, :workspace_name)" + ), + { + "peer_id": LEGACY_PEER_ID, + "peer_name": LEGACY_PEER_NAME, + "workspace_name": WORKSPACE_NAME, + }, + ) + + # Create collection with the old schema (has 'name' field) + conn.execute( + text( + f'INSERT INTO "{schema}"."collections" ' + + '("id", "name", "peer_name", "workspace_name") ' + + "VALUES (:collection_id, :collection_name, :peer_name, :workspace_name)" + ), + { + "collection_id": COLLECTION_ID, + "collection_name": GLOBAL_REP_COLLECTION_NAME, + "peer_name": LEGACY_PEER_NAME, + "workspace_name": WORKSPACE_NAME, + }, + ) + + # Create document with the old schema (has 'collection_name' field) without session_name + conn.execute( + text( + f'INSERT INTO "{schema}"."documents" ' + + '("id", "collection_name", "peer_name", "workspace_name", "content") ' + + "VALUES (:document_id, :collection_name, :peer_name, :workspace_name, :content)" + ), + { + "document_id": DOCUMENT_ID, + "collection_name": GLOBAL_REP_COLLECTION_NAME, + "peer_name": LEGACY_PEER_NAME, + "workspace_name": WORKSPACE_NAME, + "content": "test content", + }, + ) + + # Bulk create many documents with the old schema (no session_name) + bulk_stmt = text( + f'INSERT INTO "{schema}"."documents" ' + + '("id", "collection_name", "peer_name", "workspace_name", "content") ' + + "VALUES (:document_id, :collection_name, :peer_name, :workspace_name, :content)" + ) + + params_template = { + "collection_name": GLOBAL_REP_COLLECTION_NAME, + "peer_name": LEGACY_PEER_NAME, + "workspace_name": WORKSPACE_NAME, + "content": "test content", + } + + total_docs = 120_000 + batch_size = 10_000 + for start in range(0, total_docs, batch_size): + end = min(start + batch_size, total_docs) + batch_params = [ + {"document_id": generate_nanoid(), **params_template} + for _ in range(start, end) + ] + + conn.execute(bulk_stmt, batch_params) + + # Speed up bulk seeding + conn.execute(text("SET LOCAL synchronous_commit = OFF")) + + # ----------------------------- + # Group 1: self-observation collections (global_representation) + # observer = peer_name, observed = peer_name + # ----------------------------- + peer_insert_stmt = text( + f'INSERT INTO "{schema}"."peers" ("id", "name", "workspace_name") ' + + "VALUES (:peer_id, :peer_name, :workspace_name)" + ) + collection_insert_stmt = text( + f'INSERT INTO "{schema}"."collections" ("id", "name", "peer_name", "workspace_name") ' + + "VALUES (:collection_id, :name, :peer_name, :workspace_name)" + ) + + total = 100_000 + batch_size = 10_000 + + for start in range(0, total, batch_size): + end = min(start + batch_size, total) + # Insert peers for group 1 + peers_params = [ + { + "peer_id": generate_nanoid(), + "peer_name": f"selfpeer_{i}", + "workspace_name": WORKSPACE_NAME, + } + for i in range(start, end) + ] + conn.execute(peer_insert_stmt, peers_params) + + # Insert collections for group 1 + collections_params = [ + { + "collection_id": generate_nanoid(), + "name": GLOBAL_REP_COLLECTION_NAME, + "peer_name": f"selfpeer_{i}", + "workspace_name": WORKSPACE_NAME, + } + for i in range(start, end) + ] + conn.execute(collection_insert_stmt, collections_params) + + # ----------------------------- + # Group 2: prefix pattern (observer_observed) + # observer = 'prefix_observer', observed extracted from name suffix + # ----------------------------- + # Create the fixed observer peer + conn.execute( + peer_insert_stmt, + { + "peer_id": generate_nanoid(), + "peer_name": "prefix_observer", + "workspace_name": WORKSPACE_NAME, + }, + ) + + for start in range(0, total, batch_size): + end = min(start + batch_size, total) + # Observed peers for group 2 + peers_params = [ + { + "peer_id": generate_nanoid(), + "peer_name": f"obs_prefix_{i}", + "workspace_name": WORKSPACE_NAME, + } + for i in range(start, end) + ] + conn.execute(peer_insert_stmt, peers_params) + + # Collections for group 2 (name = 'prefix_observer_obs_prefix_') + collections_params = [ + { + "collection_id": generate_nanoid(), + "name": f"prefix_observer_obs_prefix_{i}", + "peer_name": "prefix_observer", + "workspace_name": WORKSPACE_NAME, + } + for i in range(start, end) + ] + conn.execute(collection_insert_stmt, collections_params) + + # ----------------------------- + # Group 3: suffix pattern (observed_observer) + # observer = 'suffix_observer', observed extracted from name prefix + # ----------------------------- + # Create the fixed observer peer + conn.execute( + peer_insert_stmt, + { + "peer_id": generate_nanoid(), + "peer_name": "suffix_observer", + "workspace_name": WORKSPACE_NAME, + }, + ) + + for start in range(0, total, batch_size): + end = min(start + batch_size, total) + # Observed peers for group 3 + peers_params = [ + { + "peer_id": generate_nanoid(), + "peer_name": f"obs_suffix_{i}", + "workspace_name": WORKSPACE_NAME, + } + for i in range(start, end) + ] + conn.execute(peer_insert_stmt, peers_params) + + # Collections for group 3 (name = 'obs_suffix__suffix_observer') + collections_params = [ + { + "collection_id": generate_nanoid(), + "name": f"obs_suffix_{i}_suffix_observer", + "peer_name": "suffix_observer", + "workspace_name": WORKSPACE_NAME, + } + for i in range(start, end) + ] + conn.execute(collection_insert_stmt, collections_params) + + +@register_after_upgrade("08894082221a") +def verify_observer_observed_migration(verifier: MigrationVerifier) -> None: + """Assert that collections/documents now rely on observer/observed fields.""" + + inspector = inspect(verifier.conn) + + collection_columns = { + col["name"] + for col in inspector.get_columns("collections", schema=verifier.schema) + } + assert "observer" in collection_columns + assert "observed" in collection_columns + assert "name" not in collection_columns + + document_columns = { + col["name"] + for col in inspector.get_columns("documents", schema=verifier.schema) + } + assert "observer" in document_columns + assert "observed" in document_columns + assert "collection_name" not in document_columns + + verifier.assert_indexes_exist( + [ + ("collections", "idx_collections_observer"), + ("collections", "idx_collections_observed"), + ("documents", "idx_documents_observer"), + ("documents", "idx_documents_observed"), + ] + ) + verifier.assert_constraint_exists( + "collections", "unique_observer_observed_collection", "unique" + ) + verifier.assert_constraint_exists( + "documents", "documents_observer_observed_workspace_name_fkey", "foreign_key" + ) + + collection = verifier.conn.execute( + text( + 'SELECT "observer", "observed", "workspace_name" ' + + f'FROM "{verifier.schema}"."collections" ' + + 'WHERE "id" = :collection_id' + ), + {"collection_id": COLLECTION_ID}, + ).one() + assert collection.observer == LEGACY_PEER_NAME + assert collection.observed == EXPECTED_OBSERVED + assert collection.workspace_name == WORKSPACE_NAME + + document = verifier.conn.execute( + text( + 'SELECT "observer", "observed", "workspace_name", "session_name" ' + + f'FROM "{verifier.schema}"."documents" ' + + 'WHERE "id" = :document_id' + ), + {"document_id": DOCUMENT_ID}, + ).one() + assert document.observer == LEGACY_PEER_NAME + assert document.observed == EXPECTED_OBSERVED + assert document.workspace_name == WORKSPACE_NAME + assert document.session_name == "__global_observations__" + + # Verify no NULLs remain in documents.session_name + verifier.assert_no_nulls("documents", "session_name") + + # Verify collections mapping for Group 1 (self-observation) + count_self = verifier.conn.execute( + text( + "SELECT COUNT(1) FROM " + + f'"{verifier.schema}"."collections" ' + + 'WHERE "workspace_name" = :ws ' + + 'AND "observer" LIKE :prefix ' + + 'AND "observed" = "observer"' + ), + {"ws": WORKSPACE_NAME, "prefix": "selfpeer_%"}, + ).scalar() + assert count_self == 100_000 + + # Verify collections mapping for Group 2 (prefix pattern observer_observed) + count_prefix = verifier.conn.execute( + text( + "SELECT COUNT(1) FROM " + + f'"{verifier.schema}"."collections" ' + + 'WHERE "workspace_name" = :ws ' + + 'AND "observer" = :observer ' + + 'AND "observed" LIKE :obs_prefix' + ), + { + "ws": WORKSPACE_NAME, + "observer": "prefix_observer", + "obs_prefix": "obs_prefix_%", + }, + ).scalar() + assert count_prefix == 100_000 + + distinct_prefix_observed = verifier.conn.execute( + text( + 'SELECT COUNT(DISTINCT "observed") FROM ' + + f'"{verifier.schema}"."collections" ' + + 'WHERE "workspace_name" = :ws AND "observer" = :observer' + ), + {"ws": WORKSPACE_NAME, "observer": "prefix_observer"}, + ).scalar() + assert distinct_prefix_observed == 100_000 + + # Verify collections mapping for Group 3 (suffix pattern observed_observer) + count_suffix = verifier.conn.execute( + text( + "SELECT COUNT(1) FROM " + + f'"{verifier.schema}"."collections" ' + + 'WHERE "workspace_name" = :ws ' + + 'AND "observer" = :observer ' + + 'AND "observed" LIKE :obs_prefix' + ), + { + "ws": WORKSPACE_NAME, + "observer": "suffix_observer", + "obs_prefix": "obs_suffix_%", + }, + ).scalar() + assert count_suffix == 100_000 + + distinct_suffix_observed = verifier.conn.execute( + text( + 'SELECT COUNT(DISTINCT "observed") FROM ' + + f'"{verifier.schema}"."collections" ' + + 'WHERE "workspace_name" = :ws AND "observer" = :observer' + ), + {"ws": WORKSPACE_NAME, "observer": "suffix_observer"}, + ).scalar() + assert distinct_suffix_observed == 100_000 diff --git a/tests/alembic/revisions/test_20f89a421aff_rename_metamessage_type_to_label.py b/tests/alembic/revisions/test_20f89a421aff_rename_metamessage_type_to_label.py new file mode 100644 index 00000000..399e70f5 --- /dev/null +++ b/tests/alembic/revisions/test_20f89a421aff_rename_metamessage_type_to_label.py @@ -0,0 +1,123 @@ +"""Hooks for revision 20f89a421aff (metamessage label rename).""" + +from __future__ import annotations + +from nanoid import generate as generate_nanoid +from sqlalchemy import text + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + +APP_ID = generate_nanoid() +USER_ID = generate_nanoid() +SESSION_ID = generate_nanoid() +MESSAGE_ID = generate_nanoid() +METAMESSAGE_ID = generate_nanoid() + + +@register_before_upgrade("20f89a421aff") +def prepare_metamessage_label(verifier: MigrationVerifier) -> None: + OLD_INDEXES = ( + ("metamessages", "idx_metamessages_lookup"), + ("metamessages", "idx_metamessages_user_lookup"), + ("metamessages", "idx_metamessages_session_lookup"), + ("metamessages", "idx_metamessages_message_lookup"), + ) + verifier.assert_indexes_exist(OLD_INDEXES) + verifier.assert_constraint_exists( + "metamessages", "metamessage_type_length", "check" + ) + + schema = verifier.schema + connection = verifier.conn + inspector = verifier.get_inspector() + + connection.execute( + text( + f'INSERT INTO "{schema}"."apps" ("public_id", "name") ' + + "VALUES (:id, :name)" + ), + {"id": APP_ID, "name": "verify-app"}, + ) + + connection.execute( + text( + f'INSERT INTO "{schema}"."users" ("public_id", "name", "app_id") ' + + "VALUES (:id, :name, :app_id)" + ), + {"id": USER_ID, "name": "verify-user", "app_id": APP_ID}, + ) + + connection.execute( + text( + f'INSERT INTO "{schema}"."sessions" ' + + '("public_id", "user_id", "app_id", "is_active") ' + + "VALUES (:id, :user_id, :app_id, true)" + ), + {"id": SESSION_ID, "user_id": USER_ID, "app_id": APP_ID}, + ) + + connection.execute( + text( + f'INSERT INTO "{schema}"."messages" ' + + '("public_id", "session_id", "app_id", "user_id", "is_user", "content") ' + + "VALUES (:id, :session_id, :app_id, :user_id, true, :content)" + ), + { + "id": MESSAGE_ID, + "session_id": SESSION_ID, + "app_id": APP_ID, + "user_id": USER_ID, + "content": "legacy metamessage seed", + }, + ) + + connection.execute( + text( + f'INSERT INTO "{schema}"."metamessages" ' + + '("public_id", "metamessage_type", "content", "message_id", "user_id", "session_id", "app_id") ' + + "VALUES (:id, :type, :content, :message_id, :user_id, :session_id, :app_id)" + ), + { + "id": METAMESSAGE_ID, + "type": "seed", + "content": "legacy metamessage", + "message_id": MESSAGE_ID, + "user_id": USER_ID, + "session_id": SESSION_ID, + "app_id": APP_ID, + }, + ) + + columns = { + col["name"] for col in inspector.get_columns("metamessages", schema=schema) + } + assert "metamessage_type" in columns + assert "label" not in columns + + +@register_after_upgrade("20f89a421aff") +def verify_metamessage_label(verifier: MigrationVerifier) -> None: + NEW_INDEXES = ( + ("metamessages", "idx_metamessages_lookup"), + ("metamessages", "idx_metamessages_user_lookup"), + ("metamessages", "idx_metamessages_session_lookup"), + ("metamessages", "idx_metamessages_message_lookup"), + ) + verifier.assert_indexes_exist(NEW_INDEXES) + verifier.assert_column_exists("metamessages", "label", nullable=False) + verifier.assert_constraint_exists("metamessages", "label_length", "check") + + verifier.assert_constraint_exists( + "metamessages", "metamessage_type_length", "check", exists=False + ) + verifier.assert_column_exists("metamessages", "metamessage_type", exists=False) + + row = verifier.conn.execute( + text( + f'SELECT "label" FROM "{verifier.schema}"."metamessages" ' + + 'WHERE "public_id" = :public_id' + ), + {"public_id": METAMESSAGE_ID}, + ).one() + assert row.label == "seed" diff --git a/tests/alembic/revisions/test_556a16564f50_add_user_id_and_app_id_to_tables.py b/tests/alembic/revisions/test_556a16564f50_add_user_id_and_app_id_to_tables.py new file mode 100644 index 00000000..c0e56afa --- /dev/null +++ b/tests/alembic/revisions/test_556a16564f50_add_user_id_and_app_id_to_tables.py @@ -0,0 +1,185 @@ +"""Hooks for revision 556a16564f50 (propagate app/user identifiers).""" + +from __future__ import annotations + +from nanoid import generate as generate_nanoid +from sqlalchemy import text + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + +APP_ID = generate_nanoid() +USER_ID = generate_nanoid() +SESSION_ID = generate_nanoid() +MESSAGE_ID = generate_nanoid() +METAMESSAGE_ID = generate_nanoid() +COLLECTION_ID = generate_nanoid() +DOCUMENT_ID = generate_nanoid() + + +@register_before_upgrade("556a16564f50") +def prepare_user_app_ids(verifier: MigrationVerifier) -> None: + schema = verifier.schema + connection = verifier.conn + + # Create app + connection.execute( + text( + f'INSERT INTO "{schema}"."apps" ("public_id", "name") ' + + "VALUES (:app_id, :name)" + ), + {"app_id": APP_ID, "name": "test-app"}, + ) + + # Create user + connection.execute( + text( + f'INSERT INTO "{schema}"."users" ("public_id", "name", "app_id") ' + + "VALUES (:user_id, :name, :app_id)" + ), + {"user_id": USER_ID, "name": "test-user", "app_id": APP_ID}, + ) + + # Create session + connection.execute( + text( + f'INSERT INTO "{schema}"."sessions" ("public_id", "user_id", "is_active") ' + + "VALUES (:session_id, :user_id, true)" + ), + {"session_id": SESSION_ID, "user_id": USER_ID}, + ) + + # Create message + connection.execute( + text( + f'INSERT INTO "{schema}"."messages" ("public_id", "session_id", "is_user", "content") ' + + "VALUES (:message_id, :session_id, true, :content)" + ), + { + "message_id": MESSAGE_ID, + "session_id": SESSION_ID, + "content": "test-content", + }, + ) + + # Create metamessage + connection.execute( + text( + f'INSERT INTO "{schema}"."metamessages" ("public_id", "user_id", "content", "metamessage_type") ' + + "VALUES (:metamessage_id, :user_id, :content, :metamessage_type)" + ), + { + "metamessage_id": METAMESSAGE_ID, + "user_id": USER_ID, + "content": "test-content", + "metamessage_type": "test-type", + }, + ) + + # Create collection + connection.execute( + text( + f'INSERT INTO "{schema}"."collections" ("public_id", "name", "user_id") ' + + "VALUES (:collection_id, :name, :user_id)" + ), + { + "collection_id": COLLECTION_ID, + "name": "test-collection", + "user_id": USER_ID, + }, + ) + + # Create document + connection.execute( + text( + f'INSERT INTO "{schema}"."documents" ("public_id", "content", "collection_id") ' + + "VALUES (:document_id, :content, :collection_id)" + ), + { + "document_id": DOCUMENT_ID, + "content": "test-document-content", + "collection_id": COLLECTION_ID, + }, + ) + + +@register_after_upgrade("556a16564f50") +def verify_app_and_user_ids(verifier: MigrationVerifier) -> None: + """Ensure app/user identifiers are present after the upgrade.""" + + verifier.assert_column_exists("sessions", "app_id", nullable=False) + verifier.assert_indexes_exist([("sessions", "ix_sessions_app_id")]) + + verifier.assert_column_exists("messages", "app_id", nullable=False) + verifier.assert_column_exists("messages", "user_id", nullable=False) + verifier.assert_indexes_exist( + [ + ("messages", "ix_messages_app_id"), + ("messages", "ix_messages_user_id"), + ] + ) + + verifier.assert_column_exists("metamessages", "app_id", nullable=False) + verifier.assert_indexes_exist([("metamessages", "ix_metamessages_app_id")]) + + verifier.assert_column_exists("collections", "app_id", nullable=False) + verifier.assert_indexes_exist([("collections", "ix_collections_app_id")]) + + verifier.assert_column_exists("documents", "app_id", nullable=False) + verifier.assert_column_exists("documents", "user_id", nullable=False) + verifier.assert_indexes_exist( + [ + ("documents", "ix_documents_app_id"), + ("documents", "ix_documents_user_id"), + ] + ) + + conn = verifier.conn + schema = verifier.schema + + session_row = conn.execute( + text( + f'SELECT "app_id" FROM "{schema}"."sessions" ' + + 'WHERE "public_id" = :session_public_id' + ), + {"session_public_id": SESSION_ID}, + ).one() + assert session_row.app_id == APP_ID + + message_row = conn.execute( + text( + f'SELECT "app_id", "user_id" FROM "{schema}"."messages" ' + + 'WHERE "public_id" = :message_public_id' + ), + {"message_public_id": MESSAGE_ID}, + ).one() + assert message_row.app_id == APP_ID + assert message_row.user_id == USER_ID + + metamessage_row = conn.execute( + text( + f'SELECT "app_id" FROM "{schema}"."metamessages" ' + + 'WHERE "public_id" = :public_id' + ), + {"public_id": METAMESSAGE_ID}, + ).one() + assert metamessage_row.app_id == APP_ID + + collection_row = conn.execute( + text( + f'SELECT "app_id" FROM "{schema}"."collections" ' + + 'WHERE "public_id" = :public_id' + ), + {"public_id": COLLECTION_ID}, + ).one() + assert collection_row.app_id == APP_ID + + document_row = conn.execute( + text( + f'SELECT "app_id", "user_id" FROM "{schema}"."documents" ' + + 'WHERE "public_id" = :public_id' + ), + {"public_id": DOCUMENT_ID}, + ).one() + assert document_row.app_id == APP_ID + assert document_row.user_id == USER_ID diff --git a/tests/alembic/revisions/test_564ba40505c5_add_session_name_column_to_documents.py b/tests/alembic/revisions/test_564ba40505c5_add_session_name_column_to_documents.py new file mode 100644 index 00000000..3bcc16d9 --- /dev/null +++ b/tests/alembic/revisions/test_564ba40505c5_add_session_name_column_to_documents.py @@ -0,0 +1,211 @@ +"""Hooks for revision 564ba40505c5 (documents session_name column).""" + +from __future__ import annotations + +import json +import time + +from nanoid import generate as generate_nanoid +from sqlalchemy import text + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + +WORKSPACE_ID = generate_nanoid() +WORKSPACE_NAME = generate_nanoid() +PEER_ID = generate_nanoid() +PEER_NAME = generate_nanoid() +SESSION_ID = generate_nanoid() +SESSION_NAME = generate_nanoid() +COLLECTION_ID = generate_nanoid() +COLLECTION_NAME = generate_nanoid() +DOCUMENT_ID = generate_nanoid() + + +@register_before_upgrade("564ba40505c5") +def seed_document_session_metadata(verifier: MigrationVerifier) -> None: + schema = verifier.schema + connection = verifier.conn + verifier.assert_column_exists("documents", "session_name", exists=False) + verifier.assert_indexes_not_exist([("documents", "idx_documents_session_name")]) + verifier.assert_constraint_exists( + "documents", "fk_documents_session_workspace", "foreign_key", exists=False + ) + + # Seed workspaces, peers, and collections + connection.execute( + text( + f""" + INSERT INTO "{schema}"."workspaces" ("id", "name") + VALUES (:workspace_id, :workspace_name) + """ + ), + {"workspace_id": WORKSPACE_ID, "workspace_name": WORKSPACE_NAME}, + ) + + connection.execute( + text( + f""" + INSERT INTO "{schema}"."peers" ("id", "name", "workspace_name") + VALUES (:peer_id, :peer_name, :workspace_name) + """ + ), + { + "peer_id": PEER_ID, + "peer_name": PEER_NAME, + "workspace_name": WORKSPACE_NAME, + }, + ) + + connection.execute( + text( + f""" + INSERT INTO "{schema}"."sessions" + ("id", "name", "workspace_name", "is_active") + VALUES (:session_id, :session_name, :workspace_name, true) + """ + ), + { + "session_id": SESSION_ID, + "session_name": SESSION_NAME, + "workspace_name": WORKSPACE_NAME, + }, + ) + + # Create session_peers entry so downgrade can find user_id + connection.execute( + text( + f""" + INSERT INTO "{schema}"."session_peers" + ("workspace_name", "session_name", "peer_name") + VALUES (:workspace_name, :session_name, :peer_name) + """ + ), + { + "workspace_name": WORKSPACE_NAME, + "session_name": SESSION_NAME, + "peer_name": PEER_NAME, + }, + ) + + connection.execute( + text( + f""" + INSERT INTO "{schema}"."collections" + ("id", "name", "peer_name", "workspace_name") + VALUES (:collection_id, :collection_name, :peer_name, :workspace_name) + """ + ), + { + "collection_id": COLLECTION_ID, + "collection_name": COLLECTION_NAME, + "peer_name": PEER_NAME, + "workspace_name": WORKSPACE_NAME, + }, + ) + + # Bulk insert 500k documents with session_name in internal_metadata to test batched backfill + print("[564ba40505c5] Starting bulk insert of 500k documents...", flush=True) + t0 = time.perf_counter() + connection.execute(text("SET LOCAL synchronous_commit = OFF")) + connection.execute( + text( + f'INSERT INTO "{schema}"."documents" ' + + '("id", "collection_name", "peer_name", "workspace_name", "content", "internal_metadata") ' + + "SELECT " + + " 'bulk-' || substr(md5(random()::text), 1, 16), " + + f" '{COLLECTION_NAME}', " + + f" '{PEER_NAME}', " + + f" '{WORKSPACE_NAME}', " + + " 'seed content', " + + f" jsonb_build_object('session_name', '{SESSION_NAME}') " + + "FROM generate_series(1, :n)" + ), + {"n": 500_000}, + ) + t1 = time.perf_counter() + print(f"[564ba40505c5] Bulk insert completed in {t1 - t0:.2f}s", flush=True) + + # Add document with session_name in internal_metadata + internal_metadata = json.dumps( + { + "session_name": SESSION_NAME, + } + ) + + connection.execute( + text( + f""" + INSERT INTO "{schema}"."documents" + ( + "id", + "collection_name", + "peer_name", + "workspace_name", + "content", + "internal_metadata" + ) + VALUES ( + :document_id, + :collection_name, + :peer_name, + :workspace_name, + :content, + :internal_metadata + ) + """ + ), + { + "document_id": DOCUMENT_ID, + "collection_name": COLLECTION_NAME, + "peer_name": PEER_NAME, + "workspace_name": WORKSPACE_NAME, + "content": "seed content", + "internal_metadata": internal_metadata, + }, + ) + + +@register_after_upgrade("564ba40505c5") +def verify_document_session_column(verifier: MigrationVerifier) -> None: + """Ensure the session_name column reflects migrated metadata.""" + + verifier.assert_column_exists("documents", "session_name", nullable=True) + verifier.assert_indexes_exist([("documents", "idx_documents_session_name")]) + verifier.assert_constraint_exists( + "documents", "fk_documents_session_workspace", "foreign_key" + ) + + # Quick diagnostics: counts before assertions + total_docs = ( + verifier.conn.execute( + text(f'SELECT COUNT(*) FROM "{verifier.schema}"."documents"') + ).scalar() + or 0 + ) + null_sessions = ( + verifier.conn.execute( + text( + f'SELECT COUNT(*) FROM "{verifier.schema}"."documents" ' + + 'WHERE "session_name" IS NULL' + ) + ).scalar() + or 0 + ) + print( + f"[564ba40505c5] Documents total={total_docs}, session_name NULLs={null_sessions}", + flush=True, + ) + + # All documents that contained session_name in internal_metadata should now have session_name populated + verifier.assert_no_nulls("documents", "session_name") + + # Sanity-check seeded document preserved expected session_name + row = verifier.conn.execute( + text( + f'SELECT "session_name" FROM "{verifier.schema}"."documents" ' + + 'WHERE "id" = :document_id' + ), + {"document_id": DOCUMENT_ID}, + ).one() + assert row.session_name == SESSION_NAME diff --git a/tests/alembic/revisions/test_66e63cf2cf77_add_indexes_to_documents_table.py b/tests/alembic/revisions/test_66e63cf2cf77_add_indexes_to_documents_table.py new file mode 100644 index 00000000..524c5368 --- /dev/null +++ b/tests/alembic/revisions/test_66e63cf2cf77_add_indexes_to_documents_table.py @@ -0,0 +1,16 @@ +"""Hooks for revision 66e63cf2cf77 (documents HNSW index).""" + +from __future__ import annotations + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + + +@register_before_upgrade("66e63cf2cf77") +def prepare_documents_hnsw(verifier: MigrationVerifier) -> None: + verifier.assert_indexes_not_exist([("documents", "idx_documents_embedding_hnsw")]) + + +@register_after_upgrade("66e63cf2cf77") +def verify_documents_hnsw(verifier: MigrationVerifier) -> None: + verifier.assert_indexes_exist([("documents", "idx_documents_embedding_hnsw")]) diff --git a/tests/alembic/revisions/test_76ffba56fe8c_add_error_field_to_queueitem.py b/tests/alembic/revisions/test_76ffba56fe8c_add_error_field_to_queueitem.py new file mode 100644 index 00000000..a88fc2b6 --- /dev/null +++ b/tests/alembic/revisions/test_76ffba56fe8c_add_error_field_to_queueitem.py @@ -0,0 +1,44 @@ +"""Hooks for revision 76ffba56fe8c (queue created_at/error columns).""" + +from __future__ import annotations + +from sqlalchemy import text + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + + +@register_before_upgrade("76ffba56fe8c") +def prepare_queue_created_at(verifier: MigrationVerifier) -> None: + verifier.assert_column_exists("queue", "created_at", nullable=False, exists=False) + verifier.assert_column_exists("queue", "error", exists=False) + verifier.assert_indexes_not_exist([("queue", "ix_queue_created_at")]) + + conn = verifier.conn + schema = verifier.schema + + # Bulk insert 500k additional queue items efficiently to exercise batched backfill + conn.execute(text("SET LOCAL synchronous_commit = OFF")) + conn.execute( + text( + f'INSERT INTO "{schema}"."queue" ' + + '("session_id", "work_unit_key", "task_type", "payload", "processed") ' + + "SELECT NULL, " + + " 'seed-batch-queue-item-' || gs::text, " + + " 'representation', " + + " jsonb_build_object('marker','bulk'), " + + " false " + + "FROM generate_series(1, :n) AS gs" + ), + {"n": 500_000}, + ) + + +@register_after_upgrade("76ffba56fe8c") +def verify_queue_created_at(verifier: MigrationVerifier) -> None: + verifier.assert_column_exists("queue", "created_at", nullable=False) + verifier.assert_column_exists("queue", "error") + verifier.assert_indexes_exist([("queue", "ix_queue_created_at")]) + + # Ensure all rows have non-null created_at after backfill + verifier.assert_no_nulls("queue", "created_at") diff --git a/tests/alembic/revisions/test_88b0fb10906f_add_webhooks_table.py b/tests/alembic/revisions/test_88b0fb10906f_add_webhooks_table.py new file mode 100644 index 00000000..1b5a0148 --- /dev/null +++ b/tests/alembic/revisions/test_88b0fb10906f_add_webhooks_table.py @@ -0,0 +1,115 @@ +"""Hooks for revision 88b0fb10906f (webhooks and queue updates).""" + +from __future__ import annotations + +import json + +from nanoid import generate as generate_nanoid +from sqlalchemy import text + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + +PAYLOAD_MARKER = generate_nanoid() +WORKSPACE_NAME = generate_nanoid() +SESSION_NAME = generate_nanoid() +SENDER_NAME = generate_nanoid() +TARGET_NAME = generate_nanoid() + + +@register_before_upgrade("88b0fb10906f") +def prepare_webhooks(verifier: MigrationVerifier) -> None: + # Validate objects introduced by the migration are absent + verifier.assert_table_exists("webhook_endpoints", exists=False) + verifier.assert_column_exists("queue", "task_type", exists=False) + verifier.assert_column_exists("queue", "work_unit_key", exists=False) + + verifier.assert_column_exists( + "active_queue_sessions", "work_unit_key", exists=False + ) + verifier.assert_constraint_exists( + "active_queue_sessions", "unique_work_unit_key", "unique", exists=False + ) + + # Validate objects removed by the migration still exist + verifier.assert_column_exists("active_queue_sessions", "session_id") + verifier.assert_column_exists("active_queue_sessions", "sender_name") + verifier.assert_column_exists("active_queue_sessions", "target_name") + verifier.assert_column_exists("active_queue_sessions", "task_type") + verifier.assert_constraint_exists( + "active_queue_sessions", "unique_active_queue_session", "unique" + ) + + # Bulk insert 500k queue items to test batched backfill + verifier.conn.execute(text("SET LOCAL synchronous_commit = OFF")) + verifier.conn.execute( + text( + f'INSERT INTO "{verifier.schema}"."queue" ("payload") ' + + "SELECT jsonb_build_object(" + + " 'marker', 'bulk-' || substr(md5(random()::text), 1, 21)," + + " 'task_type', 'representation'," + + f" 'workspace_name', '{WORKSPACE_NAME}'," + + f" 'session_name', '{SESSION_NAME}'," + + f" 'sender_name', '{SENDER_NAME}'," + + f" 'target_name', '{TARGET_NAME}'" + + ") FROM generate_series(1, :n)" + ), + {"n": 500_000}, + ) + + payload = json.dumps( + { + "marker": PAYLOAD_MARKER, + "task_type": "representation", + "workspace_name": WORKSPACE_NAME, + "session_name": SESSION_NAME, + "sender_name": SENDER_NAME, + "target_name": TARGET_NAME, + } + ) + + # Add single queue item to test expected data transformation + verifier.conn.execute( + text(f'INSERT INTO "{verifier.schema}"."queue" ("payload") VALUES (:payload)'), + {"payload": payload}, + ) + + +@register_after_upgrade("88b0fb10906f") +def verify_webhooks_and_queue(verifier: MigrationVerifier) -> None: + """Validate queue enrichment and webhook table creation.""" + + verifier.assert_table_exists("webhook_endpoints") + + verifier.assert_column_exists("queue", "task_type", nullable=False) + verifier.assert_column_exists("queue", "work_unit_key", nullable=False) + verifier.assert_column_exists("active_queue_sessions", "work_unit_key") + verifier.assert_constraint_exists( + "active_queue_sessions", "unique_work_unit_key", "unique" + ) + + verifier.assert_column_exists("active_queue_sessions", "session_id", exists=False) + verifier.assert_column_exists("active_queue_sessions", "sender_name", exists=False) + verifier.assert_column_exists("active_queue_sessions", "target_name", exists=False) + verifier.assert_column_exists("active_queue_sessions", "task_type", exists=False) + verifier.assert_constraint_exists( + "active_queue_sessions", "unique_active_queue_session", "unique", exists=False + ) + + # Ensure no NULLs remain after backfill across all rows + verifier.assert_no_nulls("queue", "task_type") + verifier.assert_no_nulls("queue", "work_unit_key") + + row = verifier.conn.execute( + text( + f'SELECT "task_type", "work_unit_key" FROM "{verifier.schema}"."queue" ' + + "WHERE payload->>'marker' = :marker" + ), + {"marker": PAYLOAD_MARKER}, + ).one() + + assert row.task_type == "representation" + expected_key = ( + f"representation:{WORKSPACE_NAME}:{SESSION_NAME}:{SENDER_NAME}:{TARGET_NAME}" + ) + assert row.work_unit_key == expected_key diff --git a/tests/alembic/revisions/test_917195d9b5e9_add_messageembedding_table.py b/tests/alembic/revisions/test_917195d9b5e9_add_messageembedding_table.py new file mode 100644 index 00000000..f9c82848 --- /dev/null +++ b/tests/alembic/revisions/test_917195d9b5e9_add_messageembedding_table.py @@ -0,0 +1,26 @@ +"""Hooks for revision 917195d9b5e9 (message embeddings table).""" + +from __future__ import annotations + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + +INDEXES = ( + ("message_embeddings", "idx_message_embeddings_message_id"), + ("message_embeddings", "idx_message_embeddings_workspace_name"), + ("message_embeddings", "idx_message_embeddings_session_name"), + ("message_embeddings", "idx_message_embeddings_peer_name"), + ("message_embeddings", "idx_message_embeddings_created_at"), + ("message_embeddings", "idx_message_embeddings_embedding_hnsw"), +) + + +@register_before_upgrade("917195d9b5e9") +def prepare_message_embeddings(verifier: MigrationVerifier) -> None: + verifier.assert_table_exists("message_embeddings", exists=False) + + +@register_after_upgrade("917195d9b5e9") +def verify_message_embeddings_table(verifier: MigrationVerifier) -> None: + verifier.assert_table_exists("message_embeddings") + verifier.assert_indexes_exist(INDEXES) diff --git a/tests/alembic/revisions/test_a1b2c3d4e5f6_initial_schema.py b/tests/alembic/revisions/test_a1b2c3d4e5f6_initial_schema.py new file mode 100644 index 00000000..2dadd2fd --- /dev/null +++ b/tests/alembic/revisions/test_a1b2c3d4e5f6_initial_schema.py @@ -0,0 +1,39 @@ +"""Hooks for initial schema revision a1b2c3d4e5f6.""" + +from __future__ import annotations + +from nanoid import generate as generate_nanoid + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + +APP_ID = generate_nanoid() +APP_NAME = "seed-app" +USER_ID = generate_nanoid() +USER_NAME = "seed-user" +SESSION_ID = generate_nanoid() +SESSION_NAME = "seed-session" +MESSAGE_PUBLIC_ID = generate_nanoid() +MESSAGE_CONTENT = "seed message" +METAMESSAGE_PUBLIC_ID = generate_nanoid() +METAMESSAGE_CONTENT = "seed metamessage" +COLLECTION_ID = generate_nanoid() +COLLECTION_NAME = "seed-collection" +DOCUMENT_ID = generate_nanoid() +DOCUMENT_CONTENT = "Seed document" + + +@register_before_upgrade("a1b2c3d4e5f6") +def prepare_initial(_verifier: MigrationVerifier) -> None: + pass + + +@register_after_upgrade("a1b2c3d4e5f6") +def verify_initial_schema(verifier: MigrationVerifier) -> None: + verifier.assert_table_exists("apps", exists=True) + verifier.assert_table_exists("users", exists=True) + verifier.assert_table_exists("sessions", exists=True) + verifier.assert_table_exists("messages", exists=True) + verifier.assert_table_exists("collections", exists=True) + verifier.assert_table_exists("documents", exists=True) + verifier.assert_table_exists("metamessages", exists=True) diff --git a/tests/alembic/revisions/test_b765d82110bd_change_metamessages_to_user_level_with_.py b/tests/alembic/revisions/test_b765d82110bd_change_metamessages_to_user_level_with_.py new file mode 100644 index 00000000..524b68b4 --- /dev/null +++ b/tests/alembic/revisions/test_b765d82110bd_change_metamessages_to_user_level_with_.py @@ -0,0 +1,114 @@ +"""Hooks for revision b765d82110bd (metamessages user-level migration).""" + +from __future__ import annotations + +from nanoid import generate as generate_nanoid +from sqlalchemy import text + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + +METAMESSAGE_ID = generate_nanoid() +USER_ID = generate_nanoid() +SESSION_ID = generate_nanoid() +MESSAGE_ID = generate_nanoid() +_INDEXES = ( + ("metamessages", "idx_metamessages_user_lookup"), + ("metamessages", "idx_metamessages_session_lookup"), + ("metamessages", "idx_metamessages_message_lookup"), +) + + +@register_before_upgrade("b765d82110bd") +def prepare_metamessages(verifier: MigrationVerifier) -> None: + schema = verifier.schema + connection = verifier.conn + # Create app + APP_ID = generate_nanoid() + connection.execute( + text( + f'INSERT INTO "{schema}"."apps" ("public_id", "name") ' + + "VALUES (:app_id, :name)" + ), + {"app_id": APP_ID, "name": "test-app"}, + ) + + # Create user + connection.execute( + text( + f'INSERT INTO "{schema}"."users" ("public_id", "name", "app_id") ' + + "VALUES (:user_id, :name, :app_id)" + ), + {"user_id": USER_ID, "name": "test-user", "app_id": APP_ID}, + ) + + # Create session + connection.execute( + text( + f'INSERT INTO "{schema}"."sessions" ("public_id", "user_id", "is_active") ' + + "VALUES (:session_id, :user_id, true)" + ), + {"session_id": SESSION_ID, "user_id": USER_ID}, + ) + + # Create message + connection.execute( + text( + f'INSERT INTO "{schema}"."messages" ("public_id", "session_id", "is_user", "content") ' + + "VALUES (:message_id, :session_id, true, :content)" + ), + { + "message_id": MESSAGE_ID, + "session_id": SESSION_ID, + "content": "test-content", + }, + ) + + # Create metamessage + connection.execute( + text( + f'INSERT INTO "{schema}"."metamessages" ("public_id", "message_id", "content", "metamessage_type") ' + + "VALUES (:metamessage_id, :message_id, :content, :metamessage_type)" + ), + { + "metamessage_id": METAMESSAGE_ID, + "message_id": MESSAGE_ID, + "content": "test-content", + "metamessage_type": "test-type", + }, + ) + + +@register_after_upgrade("b765d82110bd") +def verify_metamessages_promoted(verifier: MigrationVerifier) -> None: + """Confirm metamessages rows migrated to user-level shape.""" + + verifier.assert_column_exists("metamessages", "user_id", nullable=False) + verifier.assert_column_exists("metamessages", "session_id") + verifier.assert_column_exists("metamessages", "message_id", nullable=True) + verifier.assert_indexes_exist(_INDEXES) + verifier.assert_constraint_exists( + "metamessages", "fk_metamessages_user_id_users", "foreign_key" + ) + verifier.assert_constraint_exists( + "metamessages", "fk_metamessages_session_id_sessions", "foreign_key" + ) + verifier.assert_constraint_exists( + "metamessages", "message_requires_session", "check" + ) + verifier.assert_no_nulls("metamessages", "user_id") + + row = verifier.conn.execute( + text( + 'SELECT "user_id", "session_id", "message_id", "metamessage_type", "content" ' + + f'FROM "{verifier.schema}"."metamessages" ' + + 'WHERE "public_id" = :public_id' + ), + {"public_id": METAMESSAGE_ID}, + ).one() + + assert row.user_id == USER_ID + assert row.session_id == SESSION_ID + assert row.message_id == MESSAGE_ID + assert row.metamessage_type == "test-type" + assert row.content == "test-content" diff --git a/tests/alembic/revisions/test_bb6fb3a7a643_add_message_seq_in_session_column.py b/tests/alembic/revisions/test_bb6fb3a7a643_add_message_seq_in_session_column.py new file mode 100644 index 00000000..dd83cc03 --- /dev/null +++ b/tests/alembic/revisions/test_bb6fb3a7a643_add_message_seq_in_session_column.py @@ -0,0 +1,16 @@ +"""Hooks for revision bb6fb3a7a643 (add_message_seq_in_session_column).""" + +from __future__ import annotations + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + + +@register_before_upgrade("bb6fb3a7a643") +def prepare_add_message_seq_in_session_column(_verifier: MigrationVerifier) -> None: + """Seed state and assertions before upgrading to bb6fb3a7a643.""" + + +@register_after_upgrade("bb6fb3a7a643") +def verify_add_message_seq_in_session_column(_verifier: MigrationVerifier) -> None: + """Add assertions validating the effects of bb6fb3a7a643.""" diff --git a/tests/alembic/revisions/test_c3828084f472_add_indexes_for_messages_and_.py b/tests/alembic/revisions/test_c3828084f472_add_indexes_for_messages_and_.py new file mode 100644 index 00000000..92bdb6f5 --- /dev/null +++ b/tests/alembic/revisions/test_c3828084f472_add_indexes_for_messages_and_.py @@ -0,0 +1,23 @@ +"""Hooks for revision c3828084f472 (read indexes).""" + +from __future__ import annotations + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + +_INDEXES = ( + ("users", "idx_users_app_lookup"), + ("sessions", "idx_sessions_user_lookup"), + ("messages", "idx_messages_session_lookup"), + ("metamessages", "idx_metamessages_lookup"), +) + + +@register_before_upgrade("c3828084f472") +def prepare_read_indexes(verifier: MigrationVerifier) -> None: + verifier.assert_indexes_not_exist(_INDEXES) + + +@register_after_upgrade("c3828084f472") +def verify_read_indexes(verifier: MigrationVerifier) -> None: + verifier.assert_indexes_exist(_INDEXES) diff --git a/tests/alembic/revisions/test_d429de0e5338_adopt_peer_paradigm.py b/tests/alembic/revisions/test_d429de0e5338_adopt_peer_paradigm.py new file mode 100644 index 00000000..97c72903 --- /dev/null +++ b/tests/alembic/revisions/test_d429de0e5338_adopt_peer_paradigm.py @@ -0,0 +1,286 @@ +"""Hooks for revision d429de0e5338 (adopt peer paradigm).""" + +from __future__ import annotations + +import json + +from nanoid import generate as generate_nanoid +from sqlalchemy import text + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + +APP_ID = generate_nanoid() +APP_NAME = "app-name" +USER_ID = generate_nanoid() +USER_NAME = "user-name" +SESSION_ID = generate_nanoid() +MESSAGE_ID = generate_nanoid() +MESSAGE_CONTENT = "message-content" +COLLECTION_ID = generate_nanoid() +COLLECTION_NAME = "collection-name" +DOCUMENT_ID = generate_nanoid() +QUEUE_MARKER = "seed" + + +@register_before_upgrade("d429de0e5338") +def prepare_peer_paradigm(verifier: MigrationVerifier) -> None: + """Seed legacy tables so the migration can transform real data.""" + + schema = verifier.schema + conn = verifier.conn + + conn.execute( + text( + f'INSERT INTO "{schema}"."apps" ("public_id", "name") ' + + "VALUES (:app_id, :app_name)" + ), + {"app_id": APP_ID, "app_name": APP_NAME}, + ) + + conn.execute( + text( + f'INSERT INTO "{schema}"."users" ("public_id", "name", "app_id") ' + + "VALUES (:user_id, :user_name, :app_id)" + ), + {"user_id": USER_ID, "user_name": USER_NAME, "app_id": APP_ID}, + ) + + session_result = conn.execute( + text( + f'INSERT INTO "{schema}"."sessions" ' + + '("public_id", "user_id", "app_id", "is_active") ' + + "VALUES (:session_id, :user_id, :app_id, true) RETURNING id" + ), + {"session_id": SESSION_ID, "user_id": USER_ID, "app_id": APP_ID}, + ).one() + SESSION_DB_ID = session_result.id + + conn.execute( + text( + f'INSERT INTO "{schema}"."messages" ' + + '("public_id", "session_id", "is_user", "content", "user_id", "app_id") ' + + "VALUES (:message_id, :session_id, true, :content, :user_id, :app_id)" + ), + { + "message_id": MESSAGE_ID, + "session_id": SESSION_ID, + "content": MESSAGE_CONTENT, + "user_id": USER_ID, + "app_id": APP_ID, + }, + ) + + conn.execute( + text( + f'INSERT INTO "{schema}"."collections" ' + + '("public_id", "name", "user_id", "app_id") ' + + "VALUES (:collection_id, :collection_name, :user_id, :app_id)" + ), + { + "collection_id": COLLECTION_ID, + "collection_name": COLLECTION_NAME, + "user_id": USER_ID, + "app_id": APP_ID, + }, + ) + + conn.execute( + text( + f'INSERT INTO "{schema}"."documents" ' + + '("public_id", "collection_id", "user_id", "app_id", "content") ' + + "VALUES (:document_id, :collection_id, :user_id, :app_id, :content)" + ), + { + "document_id": DOCUMENT_ID, + "collection_id": COLLECTION_ID, + "user_id": USER_ID, + "app_id": APP_ID, + "content": "document-content", + }, + ) + + queue_payload = json.dumps({"marker": QUEUE_MARKER}) + conn.execute( + text( + f'INSERT INTO "{schema}"."queue" ' + + '("session_id", "payload", "processed") ' + + "VALUES (:session_db_id, :payload, false)" + ), + { + "session_db_id": SESSION_DB_ID, + "payload": queue_payload, + }, + ) + + +@register_after_upgrade("d429de0e5338") +def verify_peer_paradigm(verifier: MigrationVerifier) -> None: + """Validate the large-scale peer paradigm migration.""" + + inspector = verifier.get_inspector() + + assert inspector.has_table("workspaces", schema=verifier.schema) + assert not inspector.has_table("apps", schema=verifier.schema) + + verifier.assert_column_exists("workspaces", "configuration", nullable=False) + verifier.assert_column_exists("workspaces", "internal_metadata", nullable=False) + verifier.assert_column_exists("peers", "workspace_name", nullable=False) + verifier.assert_column_exists("peers", "configuration", nullable=False) + verifier.assert_column_exists("sessions", "name", nullable=False) + verifier.assert_column_exists("sessions", "workspace_name", nullable=False) + verifier.assert_column_exists( + "messages", "session_name", nullable=True + ) # later changed back to required + verifier.assert_column_exists("messages", "workspace_name", nullable=False) + verifier.assert_column_exists("messages", "peer_name", nullable=False) + verifier.assert_column_exists("collections", "peer_name", nullable=False) + verifier.assert_column_exists("collections", "workspace_name", nullable=False) + verifier.assert_column_exists("collections", "internal_metadata", nullable=False) + verifier.assert_column_exists("documents", "peer_name", nullable=False) + verifier.assert_column_exists("documents", "workspace_name", nullable=False) + verifier.assert_column_exists("documents", "collection_name", nullable=False) + verifier.assert_column_exists("documents", "internal_metadata", nullable=False) + verifier.assert_column_exists("active_queue_sessions", "id", nullable=False) + verifier.assert_column_exists("active_queue_sessions", "session_id", nullable=True) + verifier.assert_column_exists("active_queue_sessions", "sender_name", nullable=True) + verifier.assert_column_exists("active_queue_sessions", "target_name", nullable=True) + verifier.assert_column_exists("active_queue_sessions", "task_type", nullable=False) + verifier.assert_column_exists("messages", "session_id", exists=False) + verifier.assert_column_exists("messages", "user_id", exists=False) + verifier.assert_column_exists("messages", "app_id", exists=False) + verifier.assert_column_exists("messages", "is_user", exists=False) + verifier.assert_indexes_exist( + [ + ("messages", "idx_messages_session_lookup"), + ("messages", "ix_messages_peer_name"), + ("messages", "ix_messages_workspace_name"), + ] + ) + + conn = verifier.conn + schema = verifier.schema + + workspace = conn.execute( + text( + 'SELECT "id", "name", "configuration", "internal_metadata" ' + + f'FROM "{schema}"."workspaces" WHERE "id" = :workspace_id' + ), + {"workspace_id": APP_ID}, + ).one() + assert workspace.name == APP_NAME + assert workspace.configuration == {} + assert workspace.internal_metadata == {} + + peer = conn.execute( + text( + 'SELECT "id", "name", "workspace_name", "configuration", "internal_metadata" ' + + f'FROM "{schema}"."peers" WHERE "id" = :peer_id' + ), + {"peer_id": USER_ID}, + ).one() + assert peer.name == USER_NAME + assert peer.workspace_name == APP_NAME + assert peer.configuration == {} + assert peer.internal_metadata == {} + + session = conn.execute( + text( + 'SELECT "id", "name", "workspace_name", "configuration", "internal_metadata" ' + + f'FROM "{schema}"."sessions" WHERE "id" = :session_id' + ), + {"session_id": SESSION_ID}, + ).one() + assert session.name == SESSION_ID + assert session.workspace_name == APP_NAME + assert session.configuration == {} + assert session.internal_metadata == {} + + session_peer = conn.execute( + text( + f'SELECT "peer_name" FROM "{schema}"."session_peers" ' + + 'WHERE "workspace_name" = :workspace_name ' + + 'AND "session_name" = :session_name AND "peer_name" = :peer_name' + ), + { + "workspace_name": APP_NAME, + "session_name": SESSION_ID, + "peer_name": USER_NAME, + }, + ).one() + assert session_peer.peer_name == USER_NAME + + message = conn.execute( + text( + 'SELECT "session_name", "workspace_name", "peer_name", "token_count", "internal_metadata" ' + + f'FROM "{schema}"."messages" WHERE "public_id" = :message_id' + ), + {"message_id": MESSAGE_ID}, + ).one() + assert message.session_name == SESSION_ID + assert message.workspace_name == APP_NAME + assert message.peer_name == USER_NAME + assert message.token_count >= 0 + assert message.internal_metadata == {} + + collection = conn.execute( + text( + 'SELECT "id", "name", "peer_name", "workspace_name", "internal_metadata" ' + + f'FROM "{schema}"."collections" WHERE "id" = :collection_id' + ), + {"collection_id": COLLECTION_ID}, + ).one() + assert collection.name == COLLECTION_NAME + assert collection.peer_name == USER_NAME + assert collection.workspace_name == APP_NAME + assert collection.internal_metadata == {} + + document = conn.execute( + text( + 'SELECT "peer_name", "workspace_name", "collection_name", "internal_metadata" ' + + f'FROM "{schema}"."documents" WHERE "id" = :document_id' + ), + {"document_id": DOCUMENT_ID}, + ).one() + assert document.peer_name == USER_NAME + assert document.workspace_name == APP_NAME + assert document.collection_name == COLLECTION_NAME + assert document.internal_metadata == {} + + queue_row = conn.execute( + text( + f'SELECT "session_id", "payload" FROM "{schema}"."queue" ' + + "WHERE payload->>'marker' = :marker" + ), + {"marker": QUEUE_MARKER}, + ).one() + assert queue_row.session_id == SESSION_ID + + verifier.assert_constraint_exists( + "messages", "fk_messages_session_name_sessions", "foreign_key" + ) + verifier.assert_constraint_exists( + "messages", "fk_messages_peer_name_peers", "foreign_key" + ) + verifier.assert_constraint_exists( + "messages", "fk_messages_workspace_name_workspaces", "foreign_key" + ) + verifier.assert_constraint_exists( + "collections", "fk_collections_peer_name_peers", "foreign_key" + ) + verifier.assert_constraint_exists( + "collections", "fk_collections_workspace_name_workspaces", "foreign_key" + ) + verifier.assert_constraint_exists( + "peers", "fk_peers_workspace_name_workspaces", "foreign_key" + ) + verifier.assert_constraint_exists("peers", "unique_name_workspace_peer", "unique") + verifier.assert_constraint_exists("sessions", "unique_session_name", "unique") + verifier.assert_constraint_exists( + "collections", "unique_name_collection_peer", "unique" + ) + verifier.assert_constraint_exists( + "active_queue_sessions", "unique_active_queue_session", "unique" + ) + + verifier.assert_table_exists("metamessages", exists=False) diff --git a/tests/alembic/scaffold.py b/tests/alembic/scaffold.py new file mode 100644 index 00000000..80cb2d59 --- /dev/null +++ b/tests/alembic/scaffold.py @@ -0,0 +1,148 @@ +"""Utility to scaffold migration hook test modules.""" + +from __future__ import annotations + +import argparse +import re +from dataclasses import dataclass +from pathlib import Path +from textwrap import dedent + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +MIGRATIONS_DIR = PROJECT_ROOT / "migrations" / "versions" +TESTS_REVISION_DIR = Path(__file__).resolve().parent / "revisions" +REVISION_INIT_PATH = TESTS_REVISION_DIR / "__init__.py" + +TEMPLATE = '''"""Hooks for revision {revision}{slug_note}.""" + +from __future__ import annotations + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + + +@register_before_upgrade("{revision}") +def prepare_{identifier}(_verifier: MigrationVerifier) -> None: + """Seed state and assertions before upgrading to {revision}.""" + + +@register_after_upgrade("{revision}") +def verify_{identifier}(_verifier: MigrationVerifier) -> None: + """Add assertions validating the effects of {revision}.""" +''' + +HEADER = '"""Register revision-specific hooks for migration verification."""' + + +@dataclass(slots=True) +class MigrationInfo: + revision: str + slug: str + path: Path + + @property + def identifier(self) -> str: + """Return a Python-safe identifier derived from the migration slug.""" + + base = re.sub(r"[^0-9a-zA-Z]+", "_", self.slug) + base = base.strip("_").lower() or "revision" + if base[0].isdigit(): + base = f"revision_{base}" + return base + + @property + def test_filename(self) -> str: + return f"test_{self.path.stem}.py" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=("Generate a revision test stub with before/after upgrade hooks.") + ) + parser.add_argument( + "revision", + help="Revision id (e.g. a1b2c3d4e5f6) or path to a migration file.", + ) + return parser.parse_args() + + +def resolve_migration(path_or_revision: str) -> MigrationInfo: + candidate = Path(path_or_revision) + if candidate.suffix == ".py" and candidate.exists(): + return MigrationInfo( + revision=candidate.stem.split("_", 1)[0], + slug=candidate.stem.split("_", 1)[1] if "_" in candidate.stem else "", + path=candidate.resolve(), + ) + + matches = sorted(MIGRATIONS_DIR.glob(f"{path_or_revision}_*.py")) + if not matches: + raise FileNotFoundError( + f"Could not find migration matching {path_or_revision!r} in {MIGRATIONS_DIR}." + ) + if len(matches) > 1: + options = ", ".join(match.name for match in matches) + raise ValueError( + f"Revision prefix {path_or_revision!r} matches multiple migrations: {options}." + + " Provide a more specific revision or path." + ) + match = matches[0] + stem = match.stem + if "_" in stem: + revision, slug = stem.split("_", 1) + else: + revision, slug = stem, "" + return MigrationInfo(revision=revision, slug=slug, path=match.resolve()) + + +def build_template(info: MigrationInfo) -> str: + slug_note = f" ({info.slug})" if info.slug else "" + content = TEMPLATE.format( + revision=info.revision, + identifier=info.identifier, + slug_note=slug_note, + ) + return dedent(content).rstrip() + "\n" + + +def write_stub(info: MigrationInfo) -> Path: + TESTS_REVISION_DIR.mkdir(parents=True, exist_ok=True) + target = TESTS_REVISION_DIR / info.test_filename + target.write_text(build_template(info), encoding="utf-8") + return target + + +def refresh_revision_init() -> None: + modules = sorted( + path.stem for path in TESTS_REVISION_DIR.glob("test_*.py") if path.is_file() + ) + import_block = "\n".join(f" {module}," for module in modules) + all_block = "\n".join(f' "{module}",' for module in modules) + new_content = ( + dedent( + f"""{HEADER} + +from . import ( +{import_block} +) + +__all__ = [ +{all_block} +] +""" + ).rstrip() + + "\n" + ) + REVISION_INIT_PATH.write_text(new_content, encoding="utf-8") + + +def main() -> None: + args = parse_args() + info = resolve_migration(args.revision) + target = write_stub(info) + refresh_revision_init() + print(f"Created {target.relative_to(PROJECT_ROOT)}") + + +if __name__ == "__main__": + main() diff --git a/tests/alembic/test_pipeline.py b/tests/alembic/test_pipeline.py new file mode 100644 index 00000000..73fba526 --- /dev/null +++ b/tests/alembic/test_pipeline.py @@ -0,0 +1,85 @@ +"""Test pipeline for running alembic migrations and corresponding test hooks in order.""" + +from __future__ import annotations + +import pytest +from alembic import command +from alembic.config import Config +from alembic.script import ScriptDirectory +from sqlalchemy import Engine + +from tests.alembic.conftest import ALEMBIC_CONFIG_PATH +from tests.alembic.registry import get_registered_hooks +from tests.alembic.verifier import MigrationVerifier + + +def _load_revision_sequence() -> tuple[str, ...]: + """Read the Alembic script directory to produce the linear revision order.""" + + script = ScriptDirectory.from_config(Config(str(ALEMBIC_CONFIG_PATH))) + revisions = list(script.walk_revisions()) # newest -> oldest + revisions.reverse() + return tuple(revision.revision for revision in revisions) + + +REVISION_SEQUENCE: tuple[str, ...] = _load_revision_sequence() +REVISION_PARAMS = [ + pytest.param(revision, id=f"{index:02d}_{revision}") + for index, revision in enumerate(REVISION_SEQUENCE, start=1) +] + + +def _test_single_revision( + revision: str, + alembic_cfg: Config, + alembic_engine: Engine, +) -> None: + """ + Test a single migration revision upgrade. + + For each revision: + - Migrate to the previous revision + - Run the before_upgrade hook to seed and validate the state of the DB before the revision + - Migrate to the current revision + - Run the after_upgrade hook to validate the state of the DB after the revision + """ + hooks_map = get_registered_hooks() + revision_order = list(REVISION_SEQUENCE) + + # Find the previous revision in the chain + revision_index = revision_order.index(revision) + previous_revision = ( + revision_order[revision_index - 1] if revision_index > 0 else "base" + ) + + # Migrate up to the previous revision using a shared connection + with alembic_engine.begin() as conn: + alembic_cfg.attributes["connection"] = conn + command.upgrade(alembic_cfg, previous_revision) + + # Run before_upgrade hook if it exists + hooks = hooks_map.get(revision) + if hooks and hooks.before_upgrade: + with alembic_engine.begin() as conn: + verifier = MigrationVerifier(conn, revision) + hooks.before_upgrade(verifier) + + # Migrate to the current revision using the same pattern + with alembic_engine.begin() as conn: + alembic_cfg.attributes["connection"] = conn + command.upgrade(alembic_cfg, revision) + + # Run after_upgrade hook if it exists + if hooks and hooks.after_upgrade: + with alembic_engine.begin() as conn: + verifier = MigrationVerifier(conn, revision) + hooks.after_upgrade(verifier) + + +@pytest.mark.parametrize("revision", REVISION_PARAMS) +def test_migration_revision( + revision: str, + alembic_cfg: Config, + alembic_engine: Engine, +) -> None: + _test_single_revision(revision, alembic_cfg, alembic_engine) diff --git a/tests/alembic/verifier.py b/tests/alembic/verifier.py new file mode 100644 index 00000000..a859f719 --- /dev/null +++ b/tests/alembic/verifier.py @@ -0,0 +1,144 @@ +"""Utilities for asserting alembic migration behaviour inside tests.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from sqlalchemy import inspect, text +from sqlalchemy.engine import Connection +from sqlalchemy.engine.reflection import Inspector + +from migrations.utils import get_schema + + +class MigrationVerifier: + """Helper to run reusable assertions against the migrated schema.""" + + def __init__(self, connection: Connection, revision: str): + self.conn: Connection = connection + self.revision: str = revision + self.schema: str = get_schema() + self._inspector: Inspector | None = None + + def assert_table_exists(self, table: str, *, exists: bool = True) -> None: + """Assert that a table exists in the schema""" + tables = self.get_inspector().get_table_names(schema=self.schema) + + if exists: + assert table in tables + else: + assert table not in tables + + def assert_column_exists( + self, + table: str, + column: str, + *, + exists: bool = True, + nullable: bool | None = None, + ) -> None: + """Assert that a column exists in the schema""" + columns = self.get_inspector().get_columns(table, schema=self.schema) + col_names = [c["name"] for c in columns] + + if exists: + assert column in col_names + else: + assert column not in col_names + + if nullable is not None: + if column not in col_names: + # Column absence was asserted above; nothing further to verify + return + + column_info = next(col for col in columns if col["name"] == column) + actual_nullable = column_info.get("nullable", True) + assert ( + actual_nullable == nullable + ), f"Column {table}.{column} nullability is {actual_nullable}; expected {nullable}" + + def assert_column_type(self, table: str, column: str, expected_type: type) -> None: + """Assert that a column has the expected type""" + columns = self.get_inspector().get_columns(table, schema=self.schema) + column_info = next((col for col in columns if col["name"] == column), None) + assert ( + column_info is not None + ), f"Column {table}.{column} not found after migration {self.revision}" + actual_type = column_info["type"] + assert isinstance( + actual_type, expected_type + ), f"Column {table}.{column} has type {type(actual_type).__name__}; expected {expected_type.__name__}" + + def assert_no_nulls(self, table: str, column: str) -> None: + """Assert that a column has no null values""" + result = self.conn.execute( + text( + f'SELECT COUNT(*) FROM "{self.schema}"."{table}" ' + + f'WHERE "{column}" IS NULL' + ) + ) + count = result.scalar() or 0 + assert ( + count == 0 + ), f"Found {count} NULL values in {table}.{column} after migration {self.revision}" + + def assert_constraint_exists( + self, + table: str, + constraint_name: str, + constraint_type: str, + *, + exists: bool = True, + ) -> None: + """Assert that a constraint exists in the schema""" + names = self.fetch_constraints(table, constraint_type) + + if exists: + assert constraint_name in names + else: + assert constraint_name not in names + + def assert_indexes_exist(self, checks: Sequence[tuple[str, str]]) -> None: + """Assert that indexes exist in the schema""" + for table_name, index_name in checks: + indexes = self.get_inspector().get_indexes(table_name, schema=self.schema) + names = [idx["name"] for idx in indexes] + assert ( + index_name in names + ), f"Index {index_name} not found on {table_name} after migration {self.revision}" + + def assert_indexes_not_exist(self, checks: Sequence[tuple[str, str]]) -> None: + """Assert that indexes do not exist in the schema""" + for table_name, index_name in checks: + indexes = self.get_inspector().get_indexes(table_name, schema=self.schema) + names = [idx["name"] for idx in indexes] + assert ( + index_name not in names + ), f"Index {index_name} still present on {table_name} after migration {self.revision}" + + def get_inspector(self) -> Inspector: + """Get the inspector for the connection""" + if self._inspector is None: + self._inspector = inspect(self.conn) + return self._inspector + + def fetch_constraints(self, table: str, constraint_type: str) -> list[str | None]: + """Collect the names of constraints in a table""" + inspector = self.get_inspector() + + if constraint_type == "unique": + constraints = inspector.get_unique_constraints(table, schema=self.schema) + elif constraint_type == "foreign_key": + constraints = inspector.get_foreign_keys(table, schema=self.schema) + elif constraint_type == "check": + constraints = inspector.get_check_constraints(table, schema=self.schema) + elif constraint_type == "primary_key": + constraint = inspector.get_pk_constraint(table, schema=self.schema) + constraints = [constraint] if constraint else [] + else: + raise ValueError(f"Unknown constraint type: {constraint_type}") + + return [c.get("name") for c in constraints if c] + + +__all__ = ["MigrationVerifier"] diff --git a/tests/conftest.py b/tests/conftest.py index 9ba3036e..4b7cb9c1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -493,6 +493,10 @@ def mock_tracked_db(db_session: AsyncSession): patch("src.routers.sessions.tracked_db", mock_tracked_db_context), patch("src.crud.representation.tracked_db", mock_tracked_db_context), patch("src.routers.peers.tracked_db", mock_tracked_db_context), + patch("src.dreamer.dreamer.tracked_db", mock_tracked_db_context), + patch("src.dreamer.dream_scheduler.tracked_db", mock_tracked_db_context), + patch("src.dialectic.chat.tracked_db", mock_tracked_db_context), + patch("src.utils.summarizer.tracked_db", mock_tracked_db_context), ): yield diff --git a/tests/sdk/test_pagination.py b/tests/sdk/test_pagination.py index 2dcaceb9..021582cd 100644 --- a/tests/sdk/test_pagination.py +++ b/tests/sdk/test_pagination.py @@ -9,276 +9,216 @@ from sdks.python.src.honcho.peer import Peer @pytest.mark.asyncio -async def test_sync_page_get_next_page( +async def test_page_get_next_page( client_fixture: tuple[Honcho | AsyncHoncho, str], ): """ - Tests that SyncPage.get_next_page() works correctly for sync clients. + Tests that Page.get_next_page() works correctly for both sync and async clients. """ honcho_client, client_type = client_fixture - if client_type != "sync": - pytest.skip("Test only applicable to sync client") + if client_type == "async": + assert isinstance(honcho_client, AsyncHoncho) - assert isinstance(honcho_client, Honcho) + # Create multiple peers to test pagination + for i in range(15): + peer = await honcho_client.peer(id=f"pagination-test-peer-async-{i}") + await peer.get_metadata() # Create the peer - # Create multiple peers to test pagination - for i in range(15): - peer = honcho_client.peer(id=f"pagination-test-peer-{i}") - peer.get_metadata() # Create the peer + # Get first page + first_page = await honcho_client.get_peers() + assert isinstance(first_page, AsyncPage) + assert len(first_page.items) > 0 - # Get first page - SDK doesn't expose size parameter yet, so we get default pagination - first_page = honcho_client.get_peers() - assert isinstance(first_page, SyncPage) - assert len(first_page.items) > 0 + # If there's a next page, test get_next_page() + if first_page.has_next_page(): + second_page = await first_page.get_next_page() + assert second_page is not None + assert isinstance(second_page, AsyncPage) + assert second_page.page == 2 + else: + assert isinstance(honcho_client, Honcho) - # If there's a next page, test get_next_page() - if first_page.has_next_page(): - second_page = first_page.get_next_page() - assert second_page is not None - assert isinstance(second_page, SyncPage) - assert second_page.page == 2 + # Create multiple peers to test pagination + for i in range(15): + peer = honcho_client.peer(id=f"pagination-test-peer-{i}") + peer.get_metadata() # Create the peer + + # Get first page + first_page = honcho_client.get_peers() + assert isinstance(first_page, SyncPage) + assert len(first_page.items) > 0 + + # If there's a next page, test get_next_page() + if first_page.has_next_page(): + second_page = first_page.get_next_page() + assert second_page is not None + assert isinstance(second_page, SyncPage) + assert second_page.page == 2 @pytest.mark.asyncio -async def test_async_page_get_next_page( +async def test_page_transform_preserved_across_pages( client_fixture: tuple[Honcho | AsyncHoncho, str], ): """ - Tests that AsyncPage.get_next_page() works correctly for async clients. + Tests that transformation function is preserved when getting next page. """ honcho_client, client_type = client_fixture - if client_type != "async": - pytest.skip("Test only applicable to async client") + if client_type == "async": + assert isinstance(honcho_client, AsyncHoncho) - assert isinstance(honcho_client, AsyncHoncho) + # Create multiple peers to ensure pagination + for i in range(25): + peer = await honcho_client.peer(id=f"transform-test-peer-async-{i}") + await peer.get_metadata() - # Create multiple peers to test pagination - for i in range(15): - peer = await honcho_client.peer(id=f"pagination-test-peer-async-{i}") - await peer.get_metadata() # Create the peer + # Get first page + first_page = await honcho_client.get_peers() + assert isinstance(first_page, AsyncPage) - # Get first page - SDK doesn't expose size parameter yet, so we get default pagination - first_page = await honcho_client.get_peers() - assert isinstance(first_page, AsyncPage) - assert len(first_page.items) > 0 - - # If there's a next page, test get_next_page() - if first_page.has_next_page(): - second_page = await first_page.get_next_page() - assert second_page is not None - assert isinstance(second_page, AsyncPage) - assert second_page.page == 2 - - -@pytest.mark.asyncio -async def test_sync_page_transform_preserved_across_pages( - client_fixture: tuple[Honcho | AsyncHoncho, str], -): - """ - Tests that transformation function is preserved when getting next page (sync). - """ - honcho_client, client_type = client_fixture - - if client_type != "sync": - pytest.skip("Test only applicable to sync client") - - assert isinstance(honcho_client, Honcho) - - # Create multiple peers to ensure pagination - for i in range(25): - peer = honcho_client.peer(id=f"transform-test-peer-{i}") - peer.get_metadata() - - # Get first page - first_page = honcho_client.get_peers() - assert isinstance(first_page, SyncPage) - - # Verify items are Peer instances (transformed) - for item in first_page.items: - assert isinstance(item, Peer) - - # Get next page and verify transformation is preserved - if first_page.has_next_page(): - second_page = first_page.get_next_page() - assert second_page is not None - for item in second_page.items: - assert isinstance(item, Peer) - - -@pytest.mark.asyncio -async def test_async_page_transform_preserved_across_pages( - client_fixture: tuple[Honcho | AsyncHoncho, str], -): - """ - Tests that transformation function is preserved when getting next page (async). - """ - honcho_client, client_type = client_fixture - - if client_type != "async": - pytest.skip("Test only applicable to async client") - - assert isinstance(honcho_client, AsyncHoncho) - - # Create multiple peers to ensure pagination - for i in range(25): - peer = await honcho_client.peer(id=f"transform-test-peer-async-{i}") - await peer.get_metadata() - - # Get first page - first_page = await honcho_client.get_peers() - assert isinstance(first_page, AsyncPage) - - # Verify items are AsyncPeer instances (transformed) - for item in first_page.items: - assert isinstance(item, AsyncPeer) - - # Get next page and verify transformation is preserved - if first_page.has_next_page(): - second_page = await first_page.get_next_page() - assert second_page is not None - for item in second_page.items: + # Verify items are AsyncPeer instances (transformed) + for item in first_page.items: assert isinstance(item, AsyncPeer) + # Get next page and verify transformation is preserved + if first_page.has_next_page(): + second_page = await first_page.get_next_page() + assert second_page is not None + for item in second_page.items: + assert isinstance(item, AsyncPeer) + else: + assert isinstance(honcho_client, Honcho) -@pytest.mark.asyncio -async def test_sync_page_get_next_page_throws_exception_on_last_page( - client_fixture: tuple[Honcho | AsyncHoncho, str], -): - """ - Tests that get_next_page() returns None when on the last page (sync). - """ - honcho_client, client_type = client_fixture + # Create multiple peers to ensure pagination + for i in range(25): + peer = honcho_client.peer(id=f"transform-test-peer-{i}") + peer.get_metadata() - if client_type != "sync": - pytest.skip("Test only applicable to sync client") + # Get first page + first_page = honcho_client.get_peers() + assert isinstance(first_page, SyncPage) - assert isinstance(honcho_client, Honcho) + # Verify items are Peer instances (transformed) + for item in first_page.items: + assert isinstance(item, Peer) - # Create just a few peers to ensure we're on the last page - for i in range(3): - peer = honcho_client.peer(id=f"last-page-test-peer-{i}") - peer.get_metadata() - - # Get first page - first_page = honcho_client.get_peers() - assert isinstance(first_page, SyncPage) - - # Should be on last page (or only page) - if not first_page.has_next_page(): - # get_next_page should return None - try: - first_page.get_next_page() - except Exception as e: - assert isinstance(e, RuntimeError) + # Get next page and verify transformation is preserved + if first_page.has_next_page(): + second_page = first_page.get_next_page() + assert second_page is not None + for item in second_page.items: + assert isinstance(item, Peer) @pytest.mark.asyncio -async def test_async_page_get_next_page_throws_exception_on_last_page( +async def test_page_get_next_page_throws_exception_on_last_page( client_fixture: tuple[Honcho | AsyncHoncho, str], ): """ - Tests that get_next_page() returns None when on the last page (async). + Tests that get_next_page() throws RuntimeError when on the last page. """ honcho_client, client_type = client_fixture - if client_type != "async": - pytest.skip("Test only applicable to async client") + if client_type == "async": + assert isinstance(honcho_client, AsyncHoncho) - assert isinstance(honcho_client, AsyncHoncho) + # Create just a few peers to ensure we're on the last page + for i in range(3): + peer = await honcho_client.peer(id=f"last-page-test-peer-async-{i}") + await peer.get_metadata() - # Create just a few peers to ensure we're on the last page - for i in range(3): - peer = await honcho_client.peer(id=f"last-page-test-peer-async-{i}") - await peer.get_metadata() + # Get first page + first_page = await honcho_client.get_peers() + assert isinstance(first_page, AsyncPage) - # Get first page - first_page = await honcho_client.get_peers() - assert isinstance(first_page, AsyncPage) + # Should be on last page (or only page) + if not first_page.has_next_page(): + # get_next_page should throw runtime error + try: + await first_page.get_next_page() + except Exception as e: + assert isinstance(e, RuntimeError) + else: + assert isinstance(honcho_client, Honcho) - # Should be on last page (or only page) - if not first_page.has_next_page(): - # get_next_page should throw runtime error - try: - await first_page.get_next_page() - except Exception as e: - assert isinstance(e, RuntimeError) + # Create just a few peers to ensure we're on the last page + for i in range(3): + peer = honcho_client.peer(id=f"last-page-test-peer-{i}") + peer.get_metadata() + + # Get first page + first_page = honcho_client.get_peers() + assert isinstance(first_page, SyncPage) + + # Should be on last page (or only page) + if not first_page.has_next_page(): + # get_next_page should throw runtime error + try: + first_page.get_next_page() + except Exception as e: + assert isinstance(e, RuntimeError) @pytest.mark.asyncio -async def test_sync_page_manual_pagination( +async def test_page_manual_pagination( client_fixture: tuple[Honcho | AsyncHoncho, str], ): """ - Tests manual pagination with get_next_page (sync). + Tests manual pagination with get_next_page. """ honcho_client, client_type = client_fixture - if client_type != "sync": - pytest.skip("Test only applicable to sync client") + if client_type == "async": + assert isinstance(honcho_client, AsyncHoncho) - assert isinstance(honcho_client, Honcho) + # Create enough peers to ensure multiple pages + for i in range(25): + peer = await honcho_client.peer(id=f"manual-pagination-test-peer-async-{i}") + await peer.get_metadata() - # Create enough peers to ensure multiple pages - for i in range(25): - peer = honcho_client.peer(id=f"manual-pagination-test-peer-{i}") - peer.get_metadata() + # Collect items via manual pagination + manual_items = [] + manual_page = await honcho_client.get_peers() + page_count = 0 - # Collect items via manual pagination - manual_items = [] - manual_page = honcho_client.get_peers() - page_count = 0 + while manual_page is not None: + page_count += 1 + manual_items.extend(manual_page.items) # pyright: ignore - while manual_page is not None: - page_count += 1 - manual_items.extend(manual_page.items) # pyright: ignore + if not manual_page.has_next_page(): + break - if not manual_page.has_next_page(): - break + manual_page = await manual_page.get_next_page() - manual_page = manual_page.get_next_page() + # Should have collected all items + assert len(manual_items) >= 25 # pyright: ignore + # All items should be AsyncPeer instances + assert all(isinstance(item, AsyncPeer) for item in manual_items) # pyright: ignore + else: + assert isinstance(honcho_client, Honcho) - # Should have collected all items - assert len(manual_items) >= 25 # pyright: ignore - # All items should be Peer instances - assert all(isinstance(item, Peer) for item in manual_items) # pyright: ignore + # Create enough peers to ensure multiple pages + for i in range(25): + peer = honcho_client.peer(id=f"manual-pagination-test-peer-{i}") + peer.get_metadata() + # Collect items via manual pagination + manual_items = [] + manual_page = honcho_client.get_peers() + page_count = 0 -@pytest.mark.asyncio -async def test_async_page_manual_pagination( - client_fixture: tuple[Honcho | AsyncHoncho, str], -): - """ - Tests manual pagination with get_next_page (async). - """ - honcho_client, client_type = client_fixture + while manual_page is not None: + page_count += 1 + manual_items.extend(manual_page.items) # pyright: ignore - if client_type != "async": - pytest.skip("Test only applicable to async client") + if not manual_page.has_next_page(): + break - assert isinstance(honcho_client, AsyncHoncho) + manual_page = manual_page.get_next_page() - # Create enough peers to ensure multiple pages - for i in range(25): - peer = await honcho_client.peer(id=f"manual-pagination-test-peer-async-{i}") - await peer.get_metadata() - - # Collect items via manual pagination - manual_items = [] - manual_page = await honcho_client.get_peers() - page_count = 0 - - while manual_page is not None: - page_count += 1 - manual_items.extend(manual_page.items) # pyright: ignore - - if not manual_page.has_next_page(): - break - - manual_page = await manual_page.get_next_page() - - # Should have collected all items - assert len(manual_items) >= 25 # pyright: ignore - # All items should be AsyncPeer instances - assert all(isinstance(item, AsyncPeer) for item in manual_items) # pyright: ignore + # Should have collected all items + assert len(manual_items) >= 25 # pyright: ignore + # All items should be Peer instances + assert all(isinstance(item, Peer) for item in manual_items) # pyright: ignore