add typing to tests

Author: vgvoleg

tests/query/test_query_session.py

@@ -1,20 +1,22 @@
 import pytest
 
+from ydb.query.session import QuerySessionSync
 
-def _check_session_state_empty(session):
+
+def _check_session_state_empty(session: QuerySessionSync):
     assert session._state.session_id is None
     assert session._state.node_id is None
     assert not session._state.attached
 
 
-def _check_session_state_full(session):
+def _check_session_state_full(session: QuerySessionSync):
     assert session._state.session_id is not None
     assert session._state.node_id is not None
     assert session._state.attached
 
 
 class TestQuerySession:
-    def test_session_normal_lifecycle(self, session):
+    def test_session_normal_lifecycle(self, session: QuerySessionSync):
         _check_session_state_empty(session)
 
         session.create()
@@ -23,7 +25,7 @@ def test_session_normal_lifecycle(self, session):
         session.delete()
         _check_session_state_empty(session)
 
-    def test_second_create_do_nothing(self, session):
+    def test_second_create_do_nothing(self, session: QuerySessionSync):
         session.create()
         _check_session_state_full(session)
 
@@ -36,49 +38,49 @@ def test_second_create_do_nothing(self, session):
         assert session._state.session_id == session_id_before
         assert session._state.node_id == node_id_before
 
-    def test_second_delete_do_nothing(self, session):
+    def test_second_delete_do_nothing(self, session: QuerySessionSync):
         session.create()
 
         session.delete()
         session.delete()
 
-    def test_delete_before_create_not_possible(self, session):
+    def test_delete_before_create_not_possible(self, session: QuerySessionSync):
         with pytest.raises(RuntimeError):
             session.delete()
 
-    def test_create_after_delete_not_possible(self, session):
+    def test_create_after_delete_not_possible(self, session: QuerySessionSync):
         session.create()
         session.delete()
         with pytest.raises(RuntimeError):
             session.create()
 
-    def test_transaction_before_create_raises(self, session):
+    def test_transaction_before_create_raises(self, session: QuerySessionSync):
         with pytest.raises(RuntimeError):
             session.transaction()
 
-    def test_transaction_after_delete_raises(self, session):
+    def test_transaction_after_delete_raises(self, session: QuerySessionSync):
         session.create()
 
         session.delete()
 
         with pytest.raises(RuntimeError):
             session.transaction()
 
-    def test_transaction_after_create_not_raises(self, session):
+    def test_transaction_after_create_not_raises(self, session: QuerySessionSync):
         session.create()
         session.transaction()
 
-    def test_execute_before_create_raises(self, session):
+    def test_execute_before_create_raises(self, session: QuerySessionSync):
         with pytest.raises(RuntimeError):
             session.execute("select 1;")
 
-    def test_execute_after_delete_raises(self, session):
+    def test_execute_after_delete_raises(self, session: QuerySessionSync):
         session.create()
         session.delete()
         with pytest.raises(RuntimeError):
             session.execute("select 1;")
 
-    def test_basic_execute(self, session):
+    def test_basic_execute(self, session: QuerySessionSync):
         session.create()
         it = session.execute("select 1;")
         result_sets = [result_set for result_set in it]

tests/query/test_query_transaction.py

@@ -1,68 +1,69 @@
 import pytest
 
+from ydb.query.transaction import BaseTxContext
 from ydb.query.transaction import QueryTxStateEnum
 
 class TestQueryTransaction:
-    def test_tx_begin(self, tx):
+    def test_tx_begin(self, tx: BaseTxContext):
         assert tx.tx_id is None
 
         tx.begin()
         assert tx.tx_id is not None
 
-    def test_tx_allow_double_commit(self, tx):
+    def test_tx_allow_double_commit(self, tx: BaseTxContext):
         tx.begin()
         tx.commit()
         tx.commit()
 
-    def test_tx_allow_double_rollback(self, tx):
+    def test_tx_allow_double_rollback(self, tx: BaseTxContext):
         tx.begin()
         tx.rollback()
         tx.rollback()
 
-    def test_tx_commit_raises_before_begin(self, tx):
+    def test_tx_commit_raises_before_begin(self, tx: BaseTxContext):
         with pytest.raises(RuntimeError):
             tx.commit()
 
-    def test_tx_rollback_raises_before_begin(self, tx):
+    def test_tx_rollback_raises_before_begin(self, tx: BaseTxContext):
         with pytest.raises(RuntimeError):
             tx.rollback()
 
-    def test_tx_first_execute_begins_tx(self, tx):
+    def test_tx_first_execute_begins_tx(self, tx: BaseTxContext):
         tx.execute("select 1;")
         tx.commit()
 
-    def test_interactive_tx_commit(self, tx):
+    def test_interactive_tx_commit(self, tx: BaseTxContext):
         tx.execute("select 1;", commit_tx=True)
         with pytest.raises(RuntimeError):
             tx.execute("select 1;")
 
-    def test_tx_execute_raises_after_commit(self, tx):
+    def test_tx_execute_raises_after_commit(self, tx: BaseTxContext):
         tx.begin()
         tx.commit()
         with pytest.raises(RuntimeError):
             tx.execute("select 1;")
 
-    def test_tx_execute_raises_after_rollback(self, tx):
+    def test_tx_execute_raises_after_rollback(self, tx: BaseTxContext):
         tx.begin()
         tx.rollback()
         with pytest.raises(RuntimeError):
             tx.execute("select 1;")
 
-    def test_context_manager_rollbacks_tx(self, tx):
+    def test_context_manager_rollbacks_tx(self, tx: BaseTxContext):
         with tx:
             tx.begin()
 
         assert tx._tx_state._state == QueryTxStateEnum.ROLLBACKED
 
-    def test_context_manager_normal_flow(self, tx):
+    def test_context_manager_normal_flow(self, tx: BaseTxContext):
         with tx:
             tx.begin()
             tx.execute("select 1;")
             tx.commit()
 
         assert tx._tx_state._state == QueryTxStateEnum.COMMITTED
 
-    def test_context_manager_does_not_hide_exceptions(self, tx):
+    def test_context_manager_does_not_hide_exceptions(self, tx: BaseTxContext):
         with pytest.raises(RuntimeError):
             with tx:
                 tx.commit()