iceaxe 0.8.3__cp313-cp313-macosx_11_0_arm64.whl

iceaxe/__tests__/migrations/test_action_sorter.py

@@ -0,0 +1,237 @@
+import pytest
+
+from iceaxe.migrations.action_sorter import ActionTopologicalSorter
+from iceaxe.schemas.actions import DatabaseActions
+from iceaxe.schemas.db_stubs import DBObject
+
+
+class MockNode(DBObject):
+    name: str
+    table_name: str = "None"
+
+    model_config = {
+        "frozen": True,
+    }
+
+    def representation(self):
+        return f"MockNode({self.name}, {self.table_name})"
+
+    async def create(self, actor: DatabaseActions):
+        pass
+
+    async def migrate(self, previous: DBObject, actor: DatabaseActions):
+        pass
+
+    async def destroy(self, actor: DatabaseActions):
+        pass
+
+    def __hash__(self):
+        return hash(self.representation())
+
+
+def custom_topological_sort(graph_edges):
+    sorter = ActionTopologicalSorter(graph_edges)
+    sorted_objects = sorter.sort()
+    return {obj: i for i, obj in enumerate(sorted_objects)}
+
+
+def test_simple_dag():
+    A = MockNode(name="A")
+    B = MockNode(name="B")
+    C = MockNode(name="C")
+    D = MockNode(name="D")
+    graph = {D: [B, C], C: [A], B: [A], A: []}
+    result = custom_topological_sort(graph)
+    assert list(result.keys()) == [A, C, B, D]
+
+
+def test_disconnected_graph():
+    A = MockNode(name="A")
+    B = MockNode(name="B")
+    C = MockNode(name="C")
+    D = MockNode(name="D")
+    E = MockNode(name="E")
+    graph = {B: [A], A: [], D: [C], C: [], E: []}
+    result = custom_topological_sort(graph)
+    assert set(result.keys()) == {A, B, C, D, E}
+    assert result[A] < result[B]
+    assert result[C] < result[D]
+
+
+def test_single_table_grouping():
+    A = MockNode(name="A", table_name="table1")
+    B = MockNode(name="B", table_name="table1")
+    C = MockNode(name="C", table_name="table1")
+    graph = {C: [], B: [C], A: [B]}
+    result = custom_topological_sort(graph)
+    assert list(result.keys()) == [C, B, A]
+
+
+def test_multiple_table_grouping():
+    A = MockNode(name="A", table_name="table1")
+    B = MockNode(name="B", table_name="table1")
+    C = MockNode(name="C", table_name="table2")
+    D = MockNode(name="D", table_name="table2")
+    E = MockNode(name="E", table_name="table3")
+    graph = {E: [], D: [], C: [D, E], A: [C], B: [C]}
+    result = custom_topological_sort(graph)
+    assert set(result.keys()) == {A, B, C, D, E}
+    assert result[C] < result[A] and result[C] < result[B]
+    assert result[D] < result[C] and result[E] < result[C]
+
+
+def test_cross_table_references():
+    A = MockNode(name="A", table_name="table1")
+    B = MockNode(name="B", table_name="table2")
+    C = MockNode(name="C", table_name="table1")
+    D = MockNode(name="D", table_name="table2")
+    graph = {D: [], C: [D], B: [C], A: [B]}
+    result = custom_topological_sort(graph)
+    assert list(result.keys()) == [D, C, B, A]
+
+
+def test_nodes_without_table_name():
+    A = MockNode(name="A", table_name="table1")
+    B = MockNode(name="B")
+    C = MockNode(name="C", table_name="table2")
+    D = MockNode(name="D")
+    graph = {D: [], C: [D], B: [C], A: [B]}
+    result = custom_topological_sort(graph)
+    assert list(result.keys()) == [D, C, B, A]
+
+
+def test_complex_graph():
+    A = MockNode(name="A", table_name="table1")
+    B = MockNode(name="B", table_name="table1")
+    C = MockNode(name="C", table_name="table2")
+    D = MockNode(name="D", table_name="table2")
+    E = MockNode(name="E", table_name="table3")
+    F = MockNode(name="F")
+    G = MockNode(name="G", table_name="table3")
+    graph = {G: [], F: [G], E: [G], D: [F], C: [E, F], A: [C, D], B: [C]}
+    result = custom_topological_sort(graph)
+    assert set(result.keys()) == {A, B, C, D, E, F, G}
+    assert result[C] < result[A] and result[D] < result[A]
+    assert result[C] < result[B]
+    assert result[E] < result[C] and result[F] < result[C]
+    assert result[F] < result[D]
+    assert result[G] < result[E] and result[G] < result[F]
+
+
+def test_cyclic_graph():
+    A = MockNode(name="A")
+    B = MockNode(name="B")
+    C = MockNode(name="C")
+    graph = {A: [B], B: [C], C: [A]}
+    with pytest.raises(ValueError, match="Graph contains a cycle"):
+        custom_topological_sort(graph)
+
+
+def test_empty_graph():
+    graph = {}
+    result = custom_topological_sort(graph)
+    assert result == {}
+
+
+def test_single_node_graph():
+    A = MockNode(name="A")
+    graph = {A: []}
+    result = custom_topological_sort(graph)
+    assert result == {A: 0}
+
+
+def test_all_nodes_same_table():
+    A = MockNode(name="A", table_name="table1")
+    B = MockNode(name="B", table_name="table1")
+    C = MockNode(name="C", table_name="table1")
+    D = MockNode(name="D", table_name="table1")
+    graph = {D: [], B: [D], C: [D], A: [B, C]}
+    result = custom_topological_sort(graph)
+    assert list(result.keys()) == [D, B, C, A]
+
+
+def test_mixed_node_types():
+    A = MockNode(name="A", table_name="table1")
+    B = MockNode(name="B")
+    C = MockNode(name="C", table_name="table2")
+    D = MockNode(name="D", table_name="table3")
+    graph = {D: [], C: [D], B: [C], A: [B]}
+    result = custom_topological_sort(graph)
+    assert list(result.keys()) == [D, C, B, A]
+
+
+def test_large_graph_performance():
+    import random
+    import string
+    import time
+
+    def generate_large_graph(size):
+        nodes = [MockNode(name=c) for c in string.ascii_uppercase] + [
+            MockNode(name=f"N{i}") for i in range(size - 26)
+        ]
+        graph = {node: set() for node in nodes}
+        for i, node in enumerate(nodes):
+            graph[node] = set(random.sample(nodes[i + 1 :], min(5, len(nodes) - i - 1)))
+        return graph
+
+    large_graph = generate_large_graph(1000)
+    start_time = time.time()
+    result = custom_topological_sort(large_graph)
+    end_time = time.time()
+    assert len(result) == 1000
+    assert end_time - start_time < 5
+
+
+def test_graph_with_isolated_nodes():
+    A = MockNode(name="A")
+    B = MockNode(name="B")
+    C = MockNode(name="C")
+    D = MockNode(name="D")
+    E = MockNode(name="E")
+    graph = {A: [B], B: [], C: [], D: [E], E: []}
+    result = custom_topological_sort(graph)
+    assert set(result.keys()) == {A, B, C, D, E}
+    assert result[B] < result[A]
+    assert result[E] < result[D]
+
+
+@pytest.mark.parametrize(
+    "graph, expected_order",
+    [
+        (
+            {
+                MockNode(name="C"): set(),
+                MockNode(name="B"): {MockNode(name="C")},
+                MockNode(name="A"): {MockNode(name="B")},
+            },
+            ["C", "B", "A"],
+        ),
+        (
+            {
+                MockNode(name="D"): set(),
+                MockNode(name="B"): {MockNode(name="D")},
+                MockNode(name="C"): {MockNode(name="D")},
+                MockNode(name="A"): {MockNode(name="B"), MockNode(name="C")},
+            },
+            ["D", "B", "C", "A"],
+        ),
+    ],
+)
+def test_various_graph_structures(graph, expected_order):
+    result = custom_topological_sort(graph)
+    assert [node.name for node in result.keys()] == expected_order
+
+
+def test_consistent_results():
+    """
+    Test for consistent results with same input
+
+    """
+    A = MockNode(name="A", table_name="table1")
+    B = MockNode(name="B", table_name="table1")
+    C = MockNode(name="C", table_name="table2")
+    D = MockNode(name="D", table_name="table2")
+    graph = {D: [], C: [D], B: [D], A: [B, C]}
+    result1 = custom_topological_sort(graph)
+    result2 = custom_topological_sort(graph)
+    assert list(result1.keys()) == list(result2.keys())

iceaxe/__tests__/migrations/test_generator.py

@@ -0,0 +1,140 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+import pytest
+from pydantic import BaseModel
+
+from iceaxe.migrations.generator import MigrationGenerator
+from iceaxe.migrations.migration import MigrationRevisionBase
+from iceaxe.schemas.actions import ColumnType, DatabaseActions, DryRunAction
+from iceaxe.schemas.db_memory_serializer import DatabaseMemorySerializer
+from iceaxe.schemas.db_stubs import DBTable
+
+
+@pytest.mark.asyncio
+async def test_new_migration():
+    migration_generator = MigrationGenerator()
+
+    code, up_revision = await migration_generator.new_migration(
+        down_objects_with_dependencies=[
+            (
+                DBTable(
+                    table_name="test_table_a",
+                ),
+                [],
+            )
+        ],
+        up_objects_with_dependencies=[
+            (
+                DBTable(
+                    table_name="test_table_b",
+                ),
+                [],
+            )
+        ],
+        down_revision="test_down_revision",
+        user_message="test_user_message",
+    )
+
+    # Expected up
+    assert 'await migrator.actor.add_table(table_name="test_table_b")' in code
+    assert 'await migrator.actor.drop_table(table_name="test_table_a")' in code
+
+    # Expected down
+    assert 'await migrator.actor.drop_table(table_name="test_table_b")' in code
+    assert 'await migrator.actor.add_table(table_name="test_table_a")' in code
+
+    assert "Context: test_user_message" in code
+    assert 'down_revision: str | None = "test_down_revision"' in code
+
+
+def test_actions_to_code():
+    actor = DatabaseActions()
+    migration_generator = MigrationGenerator()
+
+    code = migration_generator.actions_to_code(
+        [
+            DryRunAction(
+                fn=actor.add_column,
+                kwargs={
+                    "table_name": "test_table",
+                    "column_name": "test_column",
+                    "explicit_data_type": ColumnType.VARCHAR,
+                },
+            )
+        ]
+    )
+    assert code == [
+        'await migrator.actor.add_column(table_name="test_table", column_name="test_column", explicit_data_type=ColumnType.VARCHAR)'
+    ]
+
+
+def test_actions_to_code_pass():
+    """
+    We support generating migrations where there are no schema-level changes, so users can
+    write their own data migration logic. In these cases we should pass the code-block
+    so the resulting file is still legitimate.
+
+    """
+    migration_generator = MigrationGenerator()
+    code = migration_generator.actions_to_code([])
+    assert code == ["pass"]
+
+
+class ExampleEnum(Enum):
+    A = "a"
+    B = "b"
+
+
+class ExampleModel(BaseModel):
+    value: str
+
+
+@dataclass
+class ExampleDataclass:
+    value: str
+
+
+@pytest.mark.parametrize(
+    "value, expected_value",
+    [
+        (ExampleEnum.A, "ExampleEnum.A"),
+        ("example_arg", '"example_arg"'),
+        (1, "1"),
+        (
+            {"key": "value", "nested": {"key": "value2"}},
+            '{"key": "value", "nested": {"key": "value2"}}',
+        ),
+        (
+            {
+                "key": ExampleModel(value="test"),
+                "enum": ExampleEnum.B,
+            },
+            '{"key": ExampleModel(value="test"), "enum": ExampleEnum.B}',
+        ),
+        (ExampleModel(value="test"), 'ExampleModel(value="test")'),
+        (ExampleDataclass(value="test"), 'ExampleDataclass(value="test")'),
+        (True, "True"),
+        (False, "False"),
+        (frozenset({"A", "B"}), 'frozenset({"A", "B"})'),
+        ({"A", "B"}, '{"A", "B"}'),
+        (("A",), '("A",)'),
+        (("A", "B"), '("A", "B")'),
+    ],
+)
+def test_format_arg(value: Any, expected_value: str):
+    migration_generator = MigrationGenerator()
+    assert migration_generator.format_arg(value) == expected_value
+
+
+def test_track_import():
+    migration_generator = MigrationGenerator()
+
+    migration_generator.track_import(DatabaseMemorySerializer)
+    migration_generator.track_import(MigrationRevisionBase)
+
+    assert dict(migration_generator.import_tracker) == {
+        "iceaxe.migrations.migration": {"MigrationRevisionBase"},
+        "iceaxe.schemas.db_memory_serializer": {"DatabaseMemorySerializer"},
+    }

iceaxe/__tests__/migrations/test_generics.py

@@ -0,0 +1,91 @@
+from enum import Enum
+from typing import Any, Type, Union
+
+import pytest
+
+from iceaxe.generics import (
+    _is_type_compatible,
+    remove_null_type,
+)
+
+
+class SomeEnum(Enum):
+    A = "A"
+
+
+class SomeSuperClass:
+    pass
+
+
+class SomeSubClass(SomeSuperClass):
+    pass
+
+
+@pytest.mark.parametrize(
+    "obj_type, target_type, expected",
+    [
+        # Basic types
+        (int, int, True),
+        (str, str, True),
+        (int, str, False),
+        (int, float, False),
+        # Subclasses
+        (bool, int, True),
+        (int, object, True),
+        (SomeSubClass, SomeSuperClass, True),
+        # Instance can match classes
+        (SomeSubClass(), SomeSuperClass, True),
+        ([SomeSubClass], list[SomeSuperClass], True),
+        # Enums
+        (SomeEnum, Type[Enum], True),
+        # Unions with new syntax
+        (int, Union[int, str], True),
+        (str, Union[int, str], True),
+        (float, Union[int, str], False),
+        # Unions with old syntax using type hints
+        (int, Union[int, str], True),
+        (str, Union[int, str], True),
+        (float, Union[int, str], False),
+        # Complex types involving collections
+        (list[int], list[int], True),
+        (list[int], list[str], False),
+        (dict[str, int], dict[str, int], True),
+        (dict[str, int], dict[str, str], False),
+        # More complex union cases
+        (list[int], Union[list[int], dict[str, str]], True),
+        (dict[str, str], Union[list[int], dict[str, str]], True),
+        (dict[str, float], Union[list[int], dict[str, str]], False),
+        (dict[str, str], Union[list[int], dict[str, Any]], True),
+        # Pipe operator if Python >= 3.10
+        (int, int | str, True),
+        (int | str, int | str | float, True),
+        (str, int | str, True),
+        (float, int | str, False),
+        # Nested unions
+        (int, list[int | float | str] | int | float | str, True),
+        # Optional types
+        (None, int | str | None, True),
+        (int, int | str | None, True),
+        # Value evaluation for sequences
+        ([1, 2, 3], list[int], True),
+        ([1, 2, "3"], list[int], False),
+    ],
+)
+def test_is_type_compatible(obj_type: Any, target_type: Any, expected: bool):
+    raw_result = _is_type_compatible(obj_type, target_type)
+    bool_result = raw_result != float("inf")
+    assert bool_result == expected
+
+
+@pytest.mark.parametrize(
+    "typehint, expected",
+    [
+        (int, int),
+        (str, str),
+        (int | None, int),
+        (str | None, str),
+        (Union[int, None], int),
+    ],
+)
+def test_remove_null_type(typehint: Any, expected: bool):
+    assert remove_null_type(typehint) == expected

iceaxe/__tests__/mountaineer/dependencies/test_core.py

@@ -0,0 +1,76 @@
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import asyncpg
+import pytest
+
+from iceaxe.mountaineer.config import DatabaseConfig
+from iceaxe.mountaineer.dependencies.core import get_db_connection
+from iceaxe.session import DBConnection
+
+
+@pytest.fixture(autouse=True)
+def mock_db_connect(mock_connection: AsyncMock):
+    with patch("asyncpg.connect", new_callable=AsyncMock) as mock:
+        mock.return_value = mock_connection
+
+        yield mock
+
+
+@pytest.fixture
+def mock_config():
+    return DatabaseConfig(
+        POSTGRES_HOST="test-host",
+        POSTGRES_PORT=5432,
+        POSTGRES_USER="test-user",
+        POSTGRES_PASSWORD="test-pass",
+        POSTGRES_DB="test-db",
+    )
+
+
+@pytest.fixture
+def mock_connection():
+    conn = AsyncMock(spec=asyncpg.Connection)
+    conn.close = AsyncMock()
+
+    # We need to populate the internal dsn parameters like the real query
+    conn._addr = ("test-host", 5432)
+    conn._params.user = "test-user"
+    conn._params.password = "test-pass"
+    conn._params.database = "test-db"
+
+    conn._introspect_types.return_value = (MagicMock(), MagicMock())
+
+    return conn
+
+
+@pytest.mark.asyncio
+async def test_get_db_connection_closes_after_yield(
+    mock_config: DatabaseConfig,
+    mock_connection: AsyncMock,
+    mock_db_connect: AsyncMock,
+):
+    mock_db_connect.return_value = mock_connection
+
+    # Get the generator
+    db_gen = get_db_connection(mock_config)
+
+    # Get the connection
+    connection = await anext(db_gen)  # noqa: F821
+
+    assert isinstance(connection, DBConnection)
+    assert connection.conn == mock_connection
+    mock_db_connect.assert_called_once_with(
+        host=mock_config.POSTGRES_HOST,
+        port=mock_config.POSTGRES_PORT,
+        user=mock_config.POSTGRES_USER,
+        password=mock_config.POSTGRES_PASSWORD,
+        database=mock_config.POSTGRES_DB,
+    )
+
+    # Simulate the end of the generator's scope
+    try:
+        await db_gen.aclose()
+    except StopAsyncIteration:
+        pass
+
+    mock_connection.close.assert_called_once()