Continue implementing tests. Extra nodes does not work for some reason

insani7y · insani7y · commit 9a6a22dc685b · 2024-07-22T22:39:00.000+03:00
diff --git a/psqlpy_piccolo/__pycache__/engine.cpython-312.pyc b/psqlpy_piccolo/__pycache__/engine.cpython-312.pyc
diff --git a/psqlpy_piccolo/engine.py b/psqlpy_piccolo/engine.py
@@ -276,6 +276,10 @@ def get_savepoint_id(self: Self) -> int:
         self._savepoint_id += 1
         return self._savepoint_id
 
+    async def rollback_to(self, savepoint_name: str) -> None:
+        """Use to rollback to a savepoint just using the name."""
+        await Savepoint(name=savepoint_name, transaction=self).rollback_to()
+
     async def savepoint(self: Self, name: str | None = None) -> Savepoint:
         """Create new savepoint.
 
@@ -665,7 +669,7 @@ async def run_querystring(
 
         if self.log_queries:
             self.print_query(query_id=query_id, query=querystring.__str__())
-
+        print(querystring)
         # If running inside a transaction:
         current_transaction = self.current_transaction.get()
         if current_transaction:
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -1,8 +1,65 @@
+import asyncio
+import typing
+from unittest.mock import MagicMock
+
 import pytest
 
+from tests.test_apps.mega.tables import MegaTable, SmallTable
+from tests.test_apps.music.tables import (
+    Band,
+    Concert,
+    Instrument,
+    Manager,
+    Poster,
+    RecordingStudio,
+    Shirt,
+    Ticket,
+    Venue,
+)
+
+if typing.TYPE_CHECKING:
+    from piccolo.table import Table
+
 pytestmark = [pytest.mark.anyio]
 
 
 @pytest.fixture(scope="session", autouse=True)
 def anyio_backend() -> str:
     return "asyncio"
+
+
+@pytest.fixture(autouse=True)
+async def _clean_up() -> None:
+    tables_to_clean: typing.Final[list[Table]] = [
+        MegaTable,
+        SmallTable,
+        Manager,
+        Band,
+        Venue,
+        Concert,
+        Ticket,
+        Poster,
+        Shirt,
+        RecordingStudio,
+        Instrument,
+    ]
+    for table_to_clean in tables_to_clean:
+        await table_to_clean.delete(force=True)
+
+
+class AsyncMock(MagicMock):
+    """
+    Async MagicMock for python 3.7+.
+
+    This is a workaround for the fact that MagicMock is not async compatible in
+    Python 3.7.
+    """
+
+    def __init__(self, *args: typing.Any, **kwargs: typing.Any) -> None:  # noqa: ANN401
+        super().__init__(*args, **kwargs)
+
+        # this makes asyncio.iscoroutinefunction(AsyncMock()) return True
+        self._is_coroutine = asyncio.coroutines.iscoroutine
+
+    async def __call__(self, *args: typing.Any, **kwargs: typing.Any) -> None:  # noqa: ANN401
+        return super(AsyncMock, self).__call__(*args, **kwargs)  # noqa: UP008
diff --git a/tests/test_apps/mega/piccolo_app.py b/tests/test_apps/mega/piccolo_app.py
@@ -2,14 +2,12 @@
 
 from piccolo.conf.apps import AppConfig
 
-from tests.test_apps.mega.tables import MegaTable, SmallTable
-
 CURRENT_DIRECTORY = os.path.dirname(os.path.abspath(__file__))  # noqa: PTH120, PTH100
 
 
 APP_CONFIG = AppConfig(
     app_name="mega",
-    table_classes=[MegaTable, SmallTable],
+    table_classes=[],  # During models import some circular import happening for some reason.
     migrations_folder_path=os.path.join(CURRENT_DIRECTORY, "piccolo_migrations"),  # noqa: PTH118
     commands=[],
 )
diff --git a/tests/test_extra_node.py b/tests/test_extra_node.py
@@ -0,0 +1,37 @@
+import typing
+from unittest.mock import MagicMock
+
+import pytest
+from piccolo.columns.column_types import Varchar
+from piccolo.engine import engine_finder
+from piccolo.table import Table
+
+from psqlpy_piccolo import PSQLPyEngine
+from tests.conftest import AsyncMock
+
+
+def test_extra_nodes() -> None:
+    """Make sure that other nodes can be queried."""
+    test_engine = engine_finder()
+    assert test_engine is not None
+
+    test_engine = typing.cast(PSQLPyEngine, test_engine)
+
+    EXTRA_NODE: typing.Final = MagicMock(spec=PSQLPyEngine(config=test_engine.config))  # noqa: N806
+    EXTRA_NODE.run_querystring = AsyncMock(return_value=[])
+
+    DB: typing.Final = PSQLPyEngine(  # noqa: N806
+        config=test_engine.config,
+        extra_nodes={"read_1": EXTRA_NODE},
+    )
+
+    class Manager(Table, db=DB):  # type: ignore[call-arg]
+        name = Varchar()
+
+    # Make sure the node is queried
+    Manager.select().run_sync(node="read_1")
+    assert EXTRA_NODE.run_querystring.called
+
+    # Make sure that a non existent node raises an error
+    with pytest.raises(KeyError):
+        Manager.select().run_sync(node="read_2")
diff --git a/tests/test_nested_transaction.py b/tests/test_nested_transaction.py
@@ -0,0 +1,31 @@
+import pytest
+from piccolo.engine.exceptions import TransactionError
+
+from tests.test_apps.music.tables import Manager
+
+
+async def test_nested_transactions() -> None:
+    """Make sure nested transactions databases work as expected."""
+    async with Manager._meta.db.transaction():
+        await Manager(name="Bob").save().run()
+
+        async with Manager._meta.db.transaction():
+            await Manager(name="Dave").save().run()
+
+        assert await Manager.count().run() == 2  # noqa: PLR2004
+
+    assert await Manager.count().run() == 2  # noqa: PLR2004
+
+
+async def test_nested_transactions_error() -> None:
+    """Make sure nested transactions databases work as expected."""
+    async with Manager._meta.db.transaction():
+        await Manager(name="Bob").save().run()
+
+        with pytest.raises(TransactionError):
+            async with Manager._meta.db.transaction(allow_nested=False):
+                await Manager(name="Dave").save().run()
+
+        assert await Manager.count().run() == 1
+
+    assert await Manager.count().run() == 1
diff --git a/tests/test_pool.py b/tests/test_pool.py
@@ -1,4 +1,10 @@
+import asyncio
+import os
+import tempfile
 import typing
+from unittest.mock import call, patch
+
+from piccolo.engine.sqlite import SQLiteEngine
 
 from psqlpy_piccolo import PSQLPyEngine
 from tests.test_apps.music.tables import Manager
@@ -22,3 +28,49 @@ async def test_make_query() -> None:
     assert "Bob" in [instance["name"] for instance in response]
 
     await Manager._meta.db.close_connection_pool()
+
+
+async def test_make_many_queries() -> None:
+    await Manager._meta.db.start_connection_pool()
+
+    await Manager(name="Bob").save().run()
+
+    async def get_data() -> None:
+        response = await Manager.select().run()
+        assert response[0]["name"] == "Bob"
+
+    await asyncio.gather(*[get_data() for _ in range(500)])
+
+    await Manager._meta.db.close_connection_pool()
+
+
+async def test_proxy_methods() -> None:
+    engine: typing.Final = typing.cast(PSQLPyEngine, Manager._meta.db)
+
+    # Deliberate typo ('nnn'):
+    await engine.start_connnection_pool()
+    assert engine.pool is not None
+
+    # Deliberate typo ('nnn'):
+    await engine.close_connnection_pool()
+    assert engine.pool is None
+
+
+async def test_warnings() -> None:
+    sqlite_file = os.path.join(tempfile.gettempdir(), "engine.sqlite")  # noqa: PTH118
+    engine = SQLiteEngine(path=sqlite_file)
+
+    with patch("piccolo.engine.base.colored_warning") as colored_warning:
+        await engine.start_connection_pool()
+        await engine.close_connection_pool()
+
+        assert colored_warning.call_args_list == [
+            call(
+                "Connection pooling is not supported for sqlite.",
+                stacklevel=3,
+            ),
+            call(
+                "Connection pooling is not supported for sqlite.",
+                stacklevel=3,
+            ),
+        ]
diff --git a/tests/test_transaction.py b/tests/test_transaction.py
@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+import typing
+
+import pytest
+
+from psqlpy_piccolo import PSQLPyEngine
+from tests.test_apps.music.tables import Band, Manager
+
+
+def test_atomic_error_statement() -> None:
+    """Make sure queries in a transaction aren't committed if a query fails."""
+    atomic = Band._meta.db.atomic()
+    atomic.add(
+        Band.raw("MALFORMED QUERY ... SHOULD ERROR"),
+    )
+    with pytest.raises(Exception):  # noqa: B017, PT011
+        atomic.run_sync()
+
+
+def test_atomic_succeeds_statement() -> None:
+    """Make sure that when atomic is run successfully the database is modified accordingly."""
+    atomic = Band._meta.db.atomic()
+    atomic.add(Manager.insert(Manager(name="test-manager-name")))
+    atomic.run_sync()
+    assert Manager.count().run_sync() == 1
+
+
+async def test_atomic_pool() -> None:
+    """Make sure atomic works correctly when a connection pool is active."""
+    engine = Manager._meta.db
+    await engine.start_connection_pool()
+
+    atomic = engine.atomic()
+    atomic.add(Manager.insert(Manager(name="test-manager-name")))
+
+    await atomic.run()
+    await engine.close_connection_pool()
+
+    assert Manager.count().run_sync() == 1
+
+
+async def test_transaction_error() -> None:
+    """Make sure queries in a transaction aren't committed if a query fails."""
+    with pytest.raises(Exception):  # noqa: B017, PT011, PT012
+        async with Manager._meta.db.transaction():
+            await Manager.insert(Manager(name="test-manager-name"))
+            await Manager.raw("MALFORMED QUERY ... SHOULD ERROR")
+
+    assert Manager.count().run_sync() == 0
+
+
+async def test_transaction_succeeds() -> None:
+    async with Manager._meta.db.transaction():
+        await Manager.insert(Manager(name="test-manager-name"))
+
+    assert Manager.count().run_sync() == 1
+
+
+async def test_transaction_manual_commit() -> None:
+    async with Band._meta.db.transaction() as transaction:
+        await Manager.insert(Manager(name="test-manager-name"))
+        await transaction.commit()
+
+    assert Manager.count().run_sync() == 1
+
+
+async def test_transaction_manual_rollback() -> None:
+    async with Band._meta.db.transaction() as transaction:
+        await Manager.insert(Manager(name="test-manager-name"))
+        await transaction.rollback()
+
+    assert Manager.count().run_sync() == 0
+
+
+async def test_transaction_id() -> None:
+    """An extra sanity check, that the transaction id is the same for each query inside the transaction block."""
+
+    async def get_transaction_ids() -> list[str]:
+        responses = []
+        async with Band._meta.db.transaction():
+            responses.append(await Manager.raw("SELECT txid_current()").run())
+            responses.append(await Manager.raw("SELECT txid_current()").run())
+
+        return [response[0]["txid_current"] for response in responses]
+
+    transaction_ids: typing.Final = await get_transaction_ids()
+    assert len(set(transaction_ids)) == 1
+
+    next_transaction_ids: typing.Final = await get_transaction_ids()
+    assert len(set(next_transaction_ids)) == 1
+    assert next_transaction_ids[0] != transaction_ids[0]
+
+
+async def test_transaction_exists() -> None:
+    """Make sure we can detect when code is within a transaction."""
+    engine: typing.Final = typing.cast(PSQLPyEngine, Manager._meta.db)
+
+    async with engine.transaction():
+        assert engine.transaction_exists()
+
+    assert not engine.transaction_exists()
+
+
+async def test_transaction_savepoint() -> None:
+    async with Manager._meta.db.transaction() as transaction:
+        await Manager.insert(Manager(name="Manager 1"))
+        savepoint = await transaction.savepoint()
+        await Manager.insert(Manager(name="Manager 2"))
+        await savepoint.rollback_to()
+
+        assert await Manager.select(Manager.name).run() == [{"name": "Manager 1"}]
+
+
+async def test_transaction_named_savepoint() -> None:
+    async with Manager._meta.db.transaction() as transaction:
+        await Manager.insert(Manager(name="Manager 1"))
+        await transaction.savepoint("my_savepoint1")
+        await Manager.insert(Manager(name="Manager 2"))
+        await transaction.savepoint("my_savepoint2")
+        await Manager.insert(Manager(name="Manager 3"))
+        await transaction.rollback_to("my_savepoint1")
+
+        assert await Manager.select(Manager.name).run() == [{"name": "Manager 1"}]
+
+
+async def test_savepoint_sqli_checks() -> None:
+    with pytest.raises(ValueError):  # noqa: PT011
+        async with Manager._meta.db.transaction() as transaction:
+            await transaction.savepoint(
+                "my_savepoint; SELECT * FROM Manager",
+            )
