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()