refactor cursor test

Author: vgvoleg

pyproject.toml

@@ -76,7 +76,7 @@ force-single-line = true
 
 [tool.ruff.lint.per-file-ignores]
 "**/test_*.py" = ["S", "SLF", "ANN201", "ARG", "PLR2004", "PT012"]
-"conftest.py" = ["ARG001"]
+"conftest.py" = ["S", "ARG001"]
 "__init__.py" = ["F401", "F403"]
 
 [tool.pytest.ini_options]

tests/conftest.py

@@ -185,3 +185,45 @@ def session_pool_sync(
             )
 
         yield session_pool
+
+
+@pytest.fixture
+async def session(
+    session_pool: ydb.aio.QuerySessionPool,
+) -> AsyncGenerator[ydb.aio.QuerySession]:
+    for name in ["table", "table1", "table2"]:
+        await session_pool.execute_with_retries(
+            f"""
+            DELETE FROM {name};
+            INSERT INTO {name} (id, val) VALUES
+            (0, 0),
+            (1, 1),
+            (2, 2),
+            (3, 3)
+            """
+        )
+
+    session = await session_pool.acquire()
+    yield session
+    await session_pool.release(session)
+
+
+@pytest.fixture
+def session_sync(
+    session_pool_sync: ydb.QuerySessionPool,
+) -> Generator[ydb.QuerySession]:
+    for name in ["table", "table1", "table2"]:
+        session_pool_sync.execute_with_retries(
+            f"""
+            DELETE FROM {name};
+            INSERT INTO {name} (id, val) VALUES
+            (0, 0),
+            (1, 1),
+            (2, 2),
+            (3, 3)
+            """
+        )
+
+    session = session_pool_sync.acquire()
+    yield session
+    session_pool_sync.release(session)

tests/test_connection.py

@@ -3,16 +3,17 @@
 from collections.abc import AsyncGenerator
 from collections.abc import Generator
 from contextlib import suppress
+from inspect import iscoroutine
 
 import pytest
 import pytest_asyncio
 import ydb
 import ydb_dbapi as dbapi
-from sqlalchemy.util import await_only, greenlet_spawn
-from inspect import iscoroutine
+from sqlalchemy.util import await_only
+from sqlalchemy.util import greenlet_spawn
 
 
-def maybe_await(obj):
+def maybe_await(obj: callable) -> any:
     if not iscoroutine(obj):
         return obj
     return await_only(obj)
@@ -66,7 +67,9 @@ def _test_connection(self, connection: dbapi.Connection) -> None:
         with pytest.raises(dbapi.ProgrammingError):
             maybe_await(connection.describe("/local/foo"))
 
-        maybe_await(cur.execute("CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))"))
+        maybe_await(cur.execute(
+            "CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))"
+        ))
 
         assert maybe_await(connection.check_exists("/local/foo"))
 
@@ -113,7 +116,11 @@ def _test_cursor_raw_query(self, connection: dbapi.Connection) -> None:
 
         maybe_await(cur.close())
 
-    def _test_errors(self, connection: dbapi.Connection, connect_method=dbapi.connect) -> None:
+    def _test_errors(
+        self,
+        connection: dbapi.Connection,
+        connect_method: callable = dbapi.connect
+    ) -> None:
         with pytest.raises(dbapi.InterfaceError):
             maybe_await(connect_method(
                 "localhost:2136",  # type: ignore
@@ -137,7 +144,9 @@ def _test_errors(self, connection: dbapi.Connection, connect_method=dbapi.connec
         with pytest.raises(dbapi.ProgrammingError):
             maybe_await(cur.execute("SELECT * FROM test"))
 
-        maybe_await(cur.execute("CREATE TABLE test(id Int64, PRIMARY KEY (id))"))
+        maybe_await(cur.execute(
+            "CREATE TABLE test(id Int64, PRIMARY KEY (id))"
+        ))
 
         maybe_await(cur.execute("INSERT INTO test(id) VALUES(1)"))
 
@@ -195,7 +204,7 @@ class TestAsyncConnection(BaseDBApiTestSuit):
     async def connection(
         self, connection_kwargs: dict
     ) -> AsyncGenerator[dbapi.AsyncConnection]:
-        def connect():
+        def connect() -> dbapi.AsyncConnection:
             return maybe_await(dbapi.async_connect(**connection_kwargs))
 
         conn = await greenlet_spawn(connect)