hub-core/tests/test_migration_environment.py

from __future__ import annotations

import pytest

from hub_core.migrations.roles import 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; --')