tegwick
dda93db0ac
Assistant: codex Assistant-Model: gpt-5.6-sol Assistant-Session: 01a053ff-1d6f-7fe2-ac1c-a6eb40a42a0c
51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
from unittest.mock import Mock, call
|
|
|
|
import pytest
|
|
|
|
from hub_core.migrations.roles import (
|
|
configure_migration_session,
|
|
migration_role_statement,
|
|
migration_schema_statement,
|
|
)
|
|
|
|
|
|
def test_migration_role_is_quoted_and_validated() -> None:
|
|
assert migration_role_statement("hub_runtime_owner") == 'SET ROLE "hub_runtime_owner"'
|
|
assert migration_role_statement(None) is None
|
|
with pytest.raises(ValueError, match="safe PostgreSQL role"):
|
|
migration_role_statement('owner"; DROP SCHEMA public; --')
|
|
|
|
|
|
def test_migration_schema_is_explicitly_quoted_and_validated() -> None:
|
|
assert migration_schema_statement("hub_runtime") == (
|
|
'SET search_path TO "hub_runtime", public'
|
|
)
|
|
assert migration_schema_statement(None) is None
|
|
with pytest.raises(ValueError, match="safe PostgreSQL schema"):
|
|
migration_schema_statement('hub_runtime"; DROP SCHEMA public; --')
|
|
|
|
|
|
def test_migration_session_commits_setup_before_alembic_transaction() -> None:
|
|
connection = Mock()
|
|
|
|
configure_migration_session(
|
|
connection,
|
|
role="hub_runtime_owner",
|
|
schema="hub_runtime",
|
|
)
|
|
|
|
assert connection.method_calls == [
|
|
call.exec_driver_sql('SET ROLE "hub_runtime_owner"'),
|
|
call.exec_driver_sql('SET search_path TO "hub_runtime", public'),
|
|
call.commit(),
|
|
]
|
|
|
|
|
|
def test_migration_session_does_not_open_empty_setup_transaction() -> None:
|
|
connection = Mock()
|
|
|
|
configure_migration_session(connection, role=None, schema=None)
|
|
|
|
connection.assert_not_called()
|